lib/gliner2_engine.py

"""Single-model engine: one `fastino/gliner2-multi-v1` (gliner2) replaces the
3-model zoo (GLiNER NER + MiniLM zero-shot + mREBEL relations).

`Gliner2Analyzer` subclasses `GLiNERPlusMiniLMAnalyzer` and overrides ONLY the
model-touching seams — every heuristic, DTO construction, normalization map,
cache and the public method surface (analyze_user_message / analyze_agent_message
/ analyze_key_moment / clear_session_cache) are inherited unchanged.

Seams swapped:
  * `_run_ner`            GLiNER.predict_entities → gliner2.extract_entities
  * `_run_raw_ner`        (same, for psychological discourse spans)
  * `_classify_tasks_batch`  MiniLM zero-shot argmax → gliner2.classify_text
  * `_run_classification` MiniLM multi-label scores → gliner2 binary-contrast scores
  * `extract_relations`   mREBEL → route-A curated gliner2.extract_relations (@0.5)

Relations use route A (curated 14-type schema + type-constraint/self-loop/dedup
filters); a trained relations adapter (route B) was a measured regression, so the
base model + curated schema is the shipped path. See lib/route_a_relations.py.
"""
from __future__ import annotations

import logging
from typing import Any

from lib.dto import DetectedEntity, EntityType, ExtractedRelation, RawSpan
from lib.linguistic import GLiNERPlusMiniLMAnalyzer
from lib.route_a_relations import RELS, extract_relations_route_a

logger = logging.getLogger(__name__)

GLINER2_MODEL = "fastino/gliner2-multi-v1"

# route-A relation-schema entity types -> atman EntityType, used to synthesize a
# relation endpoint's type when it can't be matched to an already-detected entity.
_ROUTE_A_TO_ATMAN: dict[str, EntityType] = {
    "person": EntityType.person,
    "organization": EntityType.organization,
    "location": EntityType.place,
    "project": EntityType.topic,
    "product": EntityType.object,
    "event": EntityType.event,
    "animal": EntityType.object,
    "health": EntityType.health_condition,
    "emotion_word": EntityType.topic,
    "profession": EntityType.skill,
}


def _clamp(x: float) -> float:
    return min(1.0, max(0.0, float(x)))


class Gliner2Analyzer(GLiNERPlusMiniLMAnalyzer):
    """GLiNERPlusMiniLMAnalyzer backed entirely by one gliner2 model."""

    def __init__(
        self,
        gliner_model: str = GLINER2_MODEL,
        device: str = "cpu",
        ner_threshold: float = 0.5,
        classification_threshold: float = 0.4,
        relation_threshold: float = 0.5,
    ) -> None:
        super().__init__(
            gliner_model=gliner_model,
            device=device,
            ner_threshold=ner_threshold,
            classification_threshold=classification_threshold,
        )
        self._relation_threshold = relation_threshold

    # ------------------------------------------------------------------
    # Lazy loader — one gliner2 model for everything
    # ------------------------------------------------------------------
    def _get_gliner(self) -> Any:
        if self._gliner is not None:
            return self._gliner
        try:
            from gliner2 import GLiNER2
        except Exception:
            logger.warning("gliner2 package not installed — analysis disabled.")
            return None
        logger.info("Loading gliner2 model %s …", self._gliner_model)
        try:
            self._gliner = GLiNER2.from_pretrained(self._gliner_model)
        except Exception:
            logger.exception("Failed to load gliner2 model %s", self._gliner_model)
            return None
        return self._gliner

    # gliner2 does classification from the same model — no separate classifier.
    def _get_classifier(self) -> Any:  # pragma: no cover - defensive
        return self._get_gliner()

    # ------------------------------------------------------------------
    # NER — entity types  (-> DetectedEntity)
    # ------------------------------------------------------------------
    def _run_ner(self, text: str) -> list[DetectedEntity]:
        if not text.strip():
            return []
        model = self._get_gliner()
        if model is None:
            return []
        sample = self._sample_text_for_models(text)
        try:
            out = model.extract_entities(
                sample,
                self._entity_type_labels(),
                threshold=self._ner_threshold,
                include_confidence=True,
                include_spans=True,
            )
        except Exception:
            logger.exception("gliner2 NER failed (len=%d)", len(text))
            return []
        ents = out.get("entities", {}) if isinstance(out, dict) else {}
        result: list[DetectedEntity] = []
        for label, items in ents.items():
            try:
                ent_type = EntityType(label)
            except (ValueError, KeyError):
                continue
            for it in items or []:
                de = self._mk_detected(it, ent_type)
                if de is not None:
                    result.append(de)
        return result

    @staticmethod
    def _mk_detected(it: Any, ent_type: EntityType) -> DetectedEntity | None:
        if isinstance(it, str):
            if not it:
                return None
            return DetectedEntity(text=it, entity_type=ent_type, confidence=1.0, span=None)
        text = it.get("text", "")
        if not text:
            return None
        conf = _clamp(it.get("confidence", it.get("score", 1.0)) or 1.0)
        span = None
        if it.get("start") is not None and it.get("end") is not None:
            span = (int(it["start"]), int(it["end"]))
        return DetectedEntity(text=text, entity_type=ent_type, confidence=conf, span=span)

    # ------------------------------------------------------------------
    # NER — raw psychological spans  (-> RawSpan, free-form labels)
    # ------------------------------------------------------------------
    def _run_raw_ner(self, text: str, labels: list[str]) -> list[RawSpan]:
        if not text.strip() or not labels:
            return []
        model = self._get_gliner()
        if model is None:
            return []
        try:
            out = model.extract_entities(
                text,
                list(labels),
                threshold=self._ner_threshold,
                include_confidence=True,
                include_spans=True,
            )
        except Exception:
            logger.exception("gliner2 raw NER failed (labels=%s)", labels[:3])
            return []
        ents = out.get("entities", {}) if isinstance(out, dict) else {}
        spans: list[RawSpan] = []
        for label, items in ents.items():
            for it in items or []:
                if isinstance(it, str):
                    if it:
                        spans.append(RawSpan(text=it, label=label, confidence=1.0, span=None))
                    continue
                t = it.get("text", "")
                if not t:
                    continue
                span = None
                if it.get("start") is not None and it.get("end") is not None:
                    span = (int(it["start"]), int(it["end"]))
                spans.append(
                    RawSpan(
                        text=t,
                        label=label,
                        confidence=_clamp(it.get("confidence", it.get("score", 0.0)) or 0.0),
                        span=span,
                    )
                )
        return spans

    # ------------------------------------------------------------------
    # Classification — argmax tasks  (-> {task: label|None})
    # ------------------------------------------------------------------
    def _classify_tasks_batch(
        self,
        text: str,
        tasks: dict[str, list[str]],
        norm_maps: dict[str, dict[str, str]] | None = None,
    ) -> dict[str, str | None]:
        norm_maps = norm_maps or {}
        if not tasks:
            return {}
        model = self._get_gliner()
        if model is None:
            return {name: None for name in tasks}
        sample = self._sample_text_for_classification(text)
        try:
            res = model.classify_text(
                sample,
                {name: list(labels) for name, labels in tasks.items()},
                threshold=0.0,
                include_confidence=True,
            )
        except Exception:
            logger.exception("gliner2 classify failed")
            return {name: None for name in tasks}
        picked: dict[str, str | None] = {}
        for name in tasks:
            r = res.get(name) if isinstance(res, dict) else None
            top: str | None = None
            if isinstance(r, dict):
                if float(r.get("confidence", 0.0) or 0.0) >= self._classification_threshold:
                    top = r.get("label")
            elif isinstance(r, str):
                top = r
            elif isinstance(r, (list, tuple)) and r:
                # format_results=False shape: [label, score]
                if len(r) < 2 or float(r[1] or 0.0) >= self._classification_threshold:
                    top = r[0]
            if top is not None and name in norm_maps:
                top = norm_maps[name].get(top, top)
            picked[name] = top
        return picked

    # ------------------------------------------------------------------
    # Classification — multi-label scores (legacy Point-K signals)
    # gliner2 classify is single-label argmax per task, so reconstruct an
    # independent per-label score via a binary contrast task per label.
    # ------------------------------------------------------------------
    def _run_classification(self, text: str, candidate_labels: list[str]) -> dict[str, float]:
        if not text.strip() or not candidate_labels:
            return {}
        model = self._get_gliner()
        if model is None:
            return {}
        sample = self._sample_text_for_classification(text)
        tasks = {f"_c{i}": [lab, "not applicable"] for i, lab in enumerate(candidate_labels)}
        try:
            res = model.classify_text(sample, tasks, threshold=0.0, include_confidence=True)
        except Exception:
            logger.exception("gliner2 multi-label classify failed")
            return {}
        scores: dict[str, float] = {}
        for i, lab in enumerate(candidate_labels):
            r = res.get(f"_c{i}") if isinstance(res, dict) else None
            if isinstance(r, dict):
                conf = float(r.get("confidence", 0.0) or 0.0)
                scores[lab] = _clamp(conf if r.get("label") == lab else 1.0 - conf)
            else:
                scores[lab] = 0.0
        return scores

    # ------------------------------------------------------------------
    # Relations — route A curated (replaces MRebelRelationExtractor)
    # ------------------------------------------------------------------
    def extract_relations(
        self, text: str, entities: list[DetectedEntity]
    ) -> list[ExtractedRelation]:
        """Curated route-A relations from the same gliner2 model.

        `entities` (already-detected DetectedEntity list) is used to resolve
        relation endpoints to typed entities; route A internally re-extracts its
        own schema-typed entities for the type-constraint filter.
        """
        model = self._get_gliner()
        if model is None or not text.strip():
            return []
        try:
            triplets = extract_relations_route_a(
                model,
                text,
                threshold=self._relation_threshold,
                ent_threshold=self._ner_threshold,
                curated=True,
            )
        except Exception:
            logger.exception("gliner2 relation extraction failed")
            return []
        by_text: dict[str, DetectedEntity] = {}
        for e in entities or []:
            by_text.setdefault(e.text.lower(), e)
        out: list[ExtractedRelation] = []
        for subj, rel, obj in triplets:
            s_ent = self._resolve_endpoint(subj, rel, "subject", by_text)
            o_ent = self._resolve_endpoint(obj, rel, "object", by_text)
            if s_ent is None or o_ent is None:
                continue
            if s_ent.text.lower() == o_ent.text.lower():
                continue
            out.append(
                ExtractedRelation(
                    subject=s_ent,
                    object=o_ent,
                    relation_type=rel,
                    confidence=1.0,
                    learned_by="gliner2-route-a",
                )
            )
        return out

    @staticmethod
    def _resolve_endpoint(
        surface: str, rel: str, role: str, by_text: dict[str, DetectedEntity]
    ) -> DetectedEntity | None:
        key = surface.lower().strip()
        if not key:
            return None
        if key in by_text:
            return by_text[key]
        for k, e in by_text.items():
            if k and (k in key or key in k):
                return e
        allow_s, allow_o = RELS.get(rel, ([], []))
        types = allow_s if role == "subject" else allow_o
        ent_type = _ROUTE_A_TO_ATMAN.get(types[0], EntityType.object) if types else EntityType.object
        try:
            return DetectedEntity(text=surface, entity_type=ent_type, confidence=1.0, span=None)
        except Exception:
            return None