handler.py

import json
from pathlib import Path
from typing import Any

import numpy as np
import onnxruntime as ort


def _stringify(value: Any) -> str:
    if value is None:
        return ""
    if isinstance(value, list):
        return ", ".join(_stringify(item) for item in value if _stringify(item))
    if isinstance(value, dict):
        return json.dumps(value, separators=(",", ":"))
    return str(value)


def _normalize_unicode_codepoint_v1(text: str) -> str:
    parts: list[str] = []
    ascii_run = ""

    def flush_ascii() -> None:
        nonlocal ascii_run
        if ascii_run:
            parts.append(ascii_run)
            ascii_run = ""

    for char in text:
        codepoint = ord(char)
        if codepoint < 128:
            ascii_run += char
        elif char.isalnum():
            flush_ascii()
            parts.append(f"u{codepoint:04x}")
        else:
            ascii_run += " "
    flush_ascii()
    return " ".join(parts)


def _field_text(property_payload: dict[str, Any]) -> str:
    semantic = property_payload.get("semantic") if isinstance(property_payload.get("semantic"), dict) else {}
    source = property_payload.get("x-mockgen-source") if isinstance(property_payload.get("x-mockgen-source"), dict) else {}
    sap_annotations = (
        source.get("sapAnnotations")
        or source.get("sourceAnnotations")
        or source.get("annotations")
        or source.get("annotationSummary")
        or source.get("sap_annotations_summary")
    )
    fields = [
        ("property_name", property_payload.get("name") or property_payload.get("property_name")),
        ("label", semantic.get("label") or property_payload.get("label")),
        ("entity_type_name", property_payload.get("entity") or property_payload.get("entity_type_name")),
        ("neighbor_properties", semantic.get("related") or property_payload.get("related") or property_payload.get("neighbor_properties")),
        ("type", property_payload.get("type")),
        ("sap_annotations_summary", sap_annotations),
        ("locale", source.get("locale") or property_payload.get("locale")),
    ]
    return " [SEP] ".join(
        f"{field}: {text_value}"
        for field, value in fields
        if (text_value := _stringify(value).strip())
    )


def _input_text(value: Any) -> str:
    if isinstance(value, str):
        stripped = value.strip()
        if stripped.startswith("{"):
            try:
                parsed = json.loads(stripped)
                if isinstance(parsed, dict):
                    return _field_text(parsed)
            except json.JSONDecodeError:
                pass
        parsed_lines = _parse_field_card(stripped)
        if parsed_lines:
            return _field_text(parsed_lines)
        return value
    if isinstance(value, dict):
        return _field_text(value)
    raise TypeError("inputs must be a classifier text string or a property metadata object")


def _parse_field_card(text: str) -> dict[str, Any] | None:
    aliases = {
        "entity": "entity",
        "entity_type": "entity",
        "entity_type_name": "entity",
        "field": "name",
        "name": "name",
        "property": "name",
        "property_name": "name",
        "label": "label",
        "type": "type",
        "related": "related",
        "neighbors": "related",
        "neighbor_properties": "related",
    }
    parsed: dict[str, Any] = {}
    for raw_line in text.splitlines():
        if ":" not in raw_line:
            continue
        key, raw_value = raw_line.split(":", 1)
        normalized_key = key.strip().lower().replace(" ", "_").replace("-", "_")
        target = aliases.get(normalized_key)
        if not target:
            continue
        value = raw_value.strip()
        if not value:
            continue
        parsed[target] = [part.strip() for part in value.split(",") if part.strip()] if target == "related" else value
    return parsed if {"entity", "name", "type"}.issubset(parsed) else None


def _softmax(scores: np.ndarray) -> np.ndarray:
    shifted = scores - np.max(scores)
    exp = np.exp(shifted)
    return exp / exp.sum()


class EndpointHandler:
    def __init__(self, path: str = ""):
        root = Path(path or ".")
        report = json.loads((root / "export-report.json").read_text(encoding="utf8"))
        self.labels = report["labels"]
        self.input_normalizer = report.get("inputNormalizer", "unicode-codepoint-v1")
        self.session = ort.InferenceSession(str(root / "classifier.onnx"), providers=["CPUExecutionProvider"])

    def __call__(self, data: dict[str, Any]) -> list[dict[str, Any]]:
        payload = data.get("inputs", data)
        text = _input_text(payload)
        if self.input_normalizer == "unicode-codepoint-v1":
            text = _normalize_unicode_codepoint_v1(text)
        outputs = self.session.run(None, {"text": np.array([[text]], dtype=object)})
        raw_scores = np.asarray(outputs[1] if len(outputs) > 1 else outputs[0]).reshape(-1).astype(float)
        scores = raw_scores if np.isclose(raw_scores.sum(), 1.0, atol=1e-3) else _softmax(raw_scores)
        ranked = sorted(
            (
                {"label": self.labels[index] if index < len(self.labels) else "unknown", "score": float(score)}
                for index, score in enumerate(scores[: len(self.labels)])
            ),
            key=lambda item: item["score"],
            reverse=True,
        )
        top_k = int(data.get("top_k", 5))
        return ranked[:top_k]