SupplyAgent/supply_infra/category_match.py

from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass
from functools import lru_cache
from threading import local
from typing import Any

import requests

from supply_infra.config import InfraSettings, get_infra_settings

_SOURCE_TYPE = "实质"
_TOP_K = 1
_MIN_SCORE = 0.8


class CategoryMatchError(RuntimeError):
    """The category matching service returned an unusable response."""


@dataclass(frozen=True)
class CategoryMatch:
    term: str
    stable_id: int
    name: str
    path: str
    description: str
    level: int
    score: float

    @property
    def reason(self) -> str:
        return (
            "category_match_api_v2: "
            f"path={self.path}; score={self.score:.4f}; stable_id={self.stable_id}"
        )


class CategoryMatchClient:
    def __init__(
        self,
        settings: InfraSettings | None = None,
        *,
        session_factory: Callable[[], requests.Session] | None = None,
    ) -> None:
        self.settings = settings or get_infra_settings()
        self._session_factory = session_factory or requests.Session
        self._thread_local = local()

    def _session(self) -> requests.Session:
        session = getattr(self._thread_local, "session", None)
        if session is None:
            session = self._session_factory()
            self._thread_local.session = session
        return session

    def match_one(
        self,
        term: str,
        *,
        description: str = "",
    ) -> CategoryMatch | None:
        normalized_term = str(term).strip()
        if not normalized_term:
            raise ValueError("term must not be empty")

        payload = {
            "source_type": _SOURCE_TYPE,
            "top_k": _TOP_K,
            "min_score": _MIN_SCORE,
            "items": [
                {
                    "term": normalized_term,
                    "description": str(description or ""),
                }
            ],
        }
        try:
            response = self._session().post(
                self.settings.category_match_api_url,
                json=payload,
                timeout=self.settings.category_match_timeout_seconds,
            )
            response.raise_for_status()
            body = response.json()
        except (requests.RequestException, ValueError) as exc:
            raise CategoryMatchError(
                f"category match request failed for term={normalized_term!r}"
            ) from exc

        return _parse_match_response(normalized_term, body)


def _parse_match_response(
    term: str,
    body: Any,
) -> CategoryMatch | None:
    if not isinstance(body, dict) or body.get("success") is not True:
        raise CategoryMatchError(f"category match response is unsuccessful for term={term!r}")

    items = body.get("items")
    if not isinstance(items, list) or len(items) != 1:
        raise CategoryMatchError(
            f"category match response must contain exactly one item for term={term!r}"
        )
    item = items[0]
    if not isinstance(item, dict) or str(item.get("term", "")).strip() != term:
        raise CategoryMatchError(f"category match response term mismatch for term={term!r}")

    matches = item.get("matches")
    if not isinstance(matches, list):
        raise CategoryMatchError(f"category match response has invalid matches for term={term!r}")
    if not matches:
        return None

    match = matches[0]
    if not isinstance(match, dict):
        raise CategoryMatchError(f"category match response has invalid match for term={term!r}")
    try:
        stable_id = int(match["stable_id"])
        score = float(match["score"])
        level = int(match["level"])
    except (KeyError, TypeError, ValueError) as exc:
        raise CategoryMatchError(
            f"category match response lacks typed fields for term={term!r}"
        ) from exc
    if score < _MIN_SCORE:
        return None

    name = str(match.get("name") or "").strip()
    path = str(match.get("path") or "").strip()
    if not name or not path:
        raise CategoryMatchError(
            f"category match response lacks name/path for term={term!r}"
        )
    return CategoryMatch(
        term=term,
        stable_id=stable_id,
        name=name,
        path=path,
        description=str(match.get("description") or ""),
        level=level,
        score=score,
    )


@lru_cache
def get_category_match_client() -> CategoryMatchClient:
    return CategoryMatchClient()