Merge branch 'features/refactoring'

Author: Paweł Kędzia

.version

@@ -0,0 +1 @@
+0.0.1

README.md

@@ -0,0 +1,23 @@
+# llm_router_services
+
+## Overview
+
+`llm_router_services` provides **HTTP services** that implement the core functionality used by the LLM‑Router’s plugin
+system.  
+The services expose guardrail and masking capabilities through Flask applications
+that can be called by the corresponding plugins in `llm_router_plugins`.
+
+Key components:
+
+| Sub‑package              | Primary purpose                                                                                                                                                                                                                    |
+|--------------------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| **guardrails/**          | Hosts the NASK‑PIB guardrail service (`nask_pib_guard_app.py`). It receives a JSON payload, chunks the text, runs a Hugging‑Face classification pipeline, and returns a safety verdict (`safe` flag + detailed per‑chunk results). |
+| **maskers/**             | Contains the **BANonymizer** (`banonymizer.py`) – a lightweight Flask service that performs token‑classification based anonymisation of input text.                                                                                |
+| **run_*.sh** scripts     | Convenience wrappers to start the services (Gunicorn for the guardrail, plain Flask for the anonymiser).                                                                                                                           |
+| **requirements‑gpu.txt** | Lists heavy dependencies (e.g., `transformers`) required for GPU‑accelerated inference.                                                                                                                                            |
+
+The services are **stateless**; they load their models once at start‑up and then serve requests over HTTP.
+
+---
+
+*Happy masking and safe routing!*

llm_router_services/__init__.py

@@ -0,0 +1 @@
+SERVICES_API_PREFIX = "/api/guardrails"

llm_router_services/guardrails/__init__.py

@@ -0,0 +1,11 @@
+from typing import Any, Dict
+from abc import ABC, abstractmethod
+
+
+class GuardrailBase(ABC):
+    """Common interface for all guardrail models."""
+
+    @abstractmethod
+    def classify_chunks(self, payload: Dict[Any, Any]) -> Dict[str, Any]:
+        """Classify the supplied payload and return a result dictionary."""
+        pass

llm_router_services/guardrails/constants.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+
+class GuardrailModelConfig(ABC):
+    """
+    Abstract base class that defines the configuration interface required by a
+    guardrail model.  Concrete implementations must provide the three fields
+    used by :class:`TextClassificationGuardrail`:
+
+    * ``pipeline_batch_size`` – size of batches sent to the HF pipeline.
+    * ``min_score_for_safe`` – threshold below which a “SAFE” label is treated as unsafe.
+    * ``min_score_for_not_safe`` – threshold above which a non‑safe label is treated as safe.
+    """
+
+    @property
+    @abstractmethod
+    def pipeline_batch_size(self) -> int: ...
+
+    @property
+    @abstractmethod
+    def min_score_for_safe(self) -> float: ...
+
+    @property
+    @abstractmethod
+    def min_score_for_not_safe(self) -> float: ...

llm_router_services/guardrails/inference/__init__.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+from typing import Any
+
+from llm_router_services.guardrails.inference.base import GuardrailBase
+from llm_router_services.guardrails.inference.config import GuardrailModelConfig
+from llm_router_services.guardrails.inference.text_classification import (
+    TextClassificationGuardrail,
+)
+
+
+def create(
+    model_type: str,
+    model_path: str,
+    device: int = -1,
+    *,
+    config: GuardrailModelConfig | None = None,
+    **kwargs: Any,
+) -> GuardrailBase:
+    """
+    Factory that builds a concrete GuardrailBase implementation.
+
+    Parameters
+    ----------
+    model_type:
+        Identifier of the concrete implementation (e.g. ``"text_classification"``).
+    model_path:
+        Path or hub identifier of the model.
+    device:
+        ``-1`` → CPU, otherwise the CUDA device index.
+    config:
+        Optional model‑specific configuration object that implements
+        :class:`GuardrailModelConfig`.  If omitted, a generic default config
+        is used.
+    kwargs:
+        Additional arguments forwarded to the concrete class.
+    """
+    if model_type == "text_classification":
+        # ``config`` may be ``None`` – the guardrail class will fall back to a
+        # generic config.
+        return TextClassificationGuardrail(
+            model_path=model_path,
+            device=device,
+            config=config,
+            **kwargs,
+        )
+    raise ValueError(f"Unsupported guardrail model_type: {model_type}")
+
+
+# Public alias expected by the Flask app
+GuardrailModelFactory = create

llm_router_services/guardrails/inference/base.py

@@ -0,0 +1,147 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, List
+
+from transformers import pipeline, AutoTokenizer, AutoConfig
+
+from llm_router_services.guardrails.inference.base import GuardrailBase
+from llm_router_services.guardrails.inference.config import GuardrailModelConfig
+from llm_router_services.guardrails.payload_handler import GuardrailPayloadExtractor
+
+
+# -----------------------------------------------------------------------
+# Default (generic) configuration – can be used when a model does not have a
+# specialized config.  It implements the GuardrailModelConfig interface.
+# -----------------------------------------------------------------------
+@dataclass(frozen=True)
+class GenericModelConfig(GuardrailModelConfig):
+    pipeline_batch_size: int = 64
+    min_score_for_safe: float = 0.5
+    min_score_for_not_safe: float = 0.5
+
+
+class TextClassificationGuardrail(GuardrailBase):
+    """
+    Generic text‑classification guardrail.
+
+    The caller supplies a concrete ``config`` object that implements
+    :class:`GuardrailModelConfig`.  This makes the guardrail reusable for any model.
+    """
+
+    def __init__(
+        self,
+        model_path: str,
+        device: int = -1,
+        max_tokens: int = 500,
+        overlap: int = 200,
+        *,
+        config: GuardrailModelConfig | None = None,
+    ):
+        # ---------------------------------------------------------------
+        # Store model‑specific thresholds & batch size
+        # ---------------------------------------------------------------
+        self._config = config or GenericModelConfig()
+
+        self._overlap = overlap
+        self._max_tokens = max_tokens
+
+        # ---------------------------------------------------------------
+        # Tokeniser & pipeline preparation (unchanged)
+        # ---------------------------------------------------------------
+        self._tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True)
+        self._model_max_length = AutoConfig.from_pretrained(
+            model_path
+        ).max_position_embeddings
+
+        if self._max_tokens > self._model_max_length:
+            self._max_tokens = self._model_max_length
+
+        self._pipeline = pipeline(
+            "text-classification",
+            model=model_path,
+            tokenizer=self._tokenizer,
+            device=device,
+            truncation=True,
+            max_length=self._max_tokens,
+        )
+
+    # -------------------------------------------------------------------
+    # Helper: convert payload → list of strings
+    # -------------------------------------------------------------------
+    @staticmethod
+    def _payload_to_string_list(payload: Dict[Any, Any]) -> List[str]:
+        try:
+            return GuardrailPayloadExtractor.extract_texts(payload)
+        except (TypeError, ValueError):
+            parts = [f"{str(k)}={str(v)}" for k, v in payload.items()]
+            return [", ".join(parts)]
+
+    # -------------------------------------------------------------------
+    # Helper: split long texts into token‑aware chunks
+    # -------------------------------------------------------------------
+    def _chunk_text(self, texts: List[str]) -> List[str]:
+        chunks: List[str] = []
+        for text in texts:
+            token_ids = self._tokenizer.encode(text, add_special_tokens=False)
+            step = self._max_tokens - self._overlap
+            for start in range(0, len(token_ids), step):
+                end = min(start + self._max_tokens, len(token_ids))
+                chunk_ids = token_ids[start:end]
+                chunk_text = self._tokenizer.decode(
+                    chunk_ids,
+                    skip_special_tokens=True,
+                    clean_up_tokenization_spaces=True,
+                )
+                chunks.append(chunk_text.strip())
+                if end == len(token_ids):
+                    break
+        return chunks
+
+    # -------------------------------------------------------------------
+    # Public API – called from the Flask endpoint
+    # -------------------------------------------------------------------
+    def classify_chunks(self, payload: Dict[Any, Any]) -> Dict[str, Any]:
+        texts = self._payload_to_string_list(payload)
+        chunks = self._chunk_text(texts=texts)
+
+        # Run inference in batches defined by the model config
+        raw_results = self._pipeline(
+            chunks, batch_size=self._config.pipeline_batch_size
+        )
+
+        # Normalise pipeline output (it can be a list of dicts or a list containing a single list)
+        flat_results = [r[0] if isinstance(r, list) else r for r in raw_results]
+
+        detailed: List[Dict[str, Any]] = []
+        for idx, (chunk, classification) in enumerate(zip(chunks, flat_results)):
+            label = classification.get("label", "")
+            score = round(classification.get("score", 0.0), 4)
+            is_safe = label.lower() == "safe"
+
+            detailed.append(
+                {
+                    "chunk_index": idx,
+                    "chunk_text": chunk,
+                    "label": label,
+                    "score": score,
+                    "safe": is_safe,
+                }
+            )
+
+        # ---------------------------------------------------------------
+        # Overall safety decision – uses the per‑model thresholds
+        # ---------------------------------------------------------------
+        overall_safe = True
+        for item in detailed:
+            if item["safe"] and item["score"] < self._config.min_score_for_safe:
+                overall_safe = False
+                break
+            if (
+                not item["safe"]
+                and item["score"] > self._config.min_score_for_not_safe
+            ):
+                overall_safe = False
+                break
+
+        return {"safe": overall_safe, "detailed": detailed}