From 4c4875fe3ae5ff9be97362bc07e67bca44937a54 Mon Sep 17 00:00:00 2001
From: crp4222 <121295348+crp4222@users.noreply.github.com>
Date: Thu, 9 Jul 2026 15:16:50 +0100
Subject: [PATCH 1/4] refactor: extract shared detector, recognizer and config
 abstractions

---
 privaite/config/schema.py              |  61 +++++-----
 privaite/middleware/limits.py          |  29 ++---
 privaite/pii/detector_base.py          | 113 +++++++++++++++++++
 privaite/pii/detector_bert_ner.py      |  68 +----------
 privaite/pii/detector_gliner.py        |  79 +++++--------
 privaite/pii/detector_mlmodel.py       |  63 +----------
 privaite/pii/recognizer_base.py        |  35 ++++++
 privaite/pii/recognizer_custom.py      |  32 +-----
 privaite/pii/recognizer_fr_date.py     |  30 +----
 privaite/pii/recognizer_location.py    |  27 +----
 tests/test_api/test_limits.py          |  68 +++++++++++
 tests/test_pii/test_windowed_detect.py | 150 +++++++++++++++++++++++++
 12 files changed, 451 insertions(+), 304 deletions(-)
 create mode 100644 privaite/pii/recognizer_base.py
 create mode 100644 tests/test_pii/test_windowed_detect.py

diff --git a/privaite/config/schema.py b/privaite/config/schema.py
index 8a15f1b..286ec8a 100644
--- a/privaite/config/schema.py
+++ b/privaite/config/schema.py
@@ -37,54 +37,47 @@ class PresidioDetectorConfig(BaseModel):
     entities: list[str] | None = None
 
 
-class MLModelDetectorConfig(BaseModel):
+# The openai/privacy-filter model's label set, shared by its ONNX and torch backends.
+_PRIVACY_FILTER_LABELS = {
+    "private_person": "PERSON",
+    "private_email": "EMAIL_ADDRESS",
+    "private_phone": "PHONE_NUMBER",
+    "private_address": "LOCATION",
+    "private_date": "DATE_TIME",
+    "private_url": "URL",
+    "account_number": "FINANCIAL",
+    "secret": "SECRET",
+}
+
+
+class _PrivacyFilterDetectorConfig(BaseModel):
+    """Shared config for the two openai/privacy-filter backends; they differ only in
+    runtime knobs (ONNX variant vs torch dtype/cache)."""
+
     enabled: bool = False
     model_name: str = "openai/privacy-filter"
     # Pin a Hugging Face revision (tag or commit sha) so a rewritten model repo
     # cannot silently swap the weights this proxy runs. None = latest.
     revision: str | None = None
-    # Off by default: a PII proxy must not execute code shipped inside a model
-    # repo unless the operator explicitly opts in for a custom-code model.
+    # Off by default: a PII proxy must not execute code shipped inside a model repo
+    # unless the operator explicitly opts in for a custom-code model.
     trust_remote_code: bool = False
     device: str = "auto"
-    torch_dtype: str = "float16"
     score_threshold: float = 0.5
+    label_mapping: dict[str, str] = Field(
+        default_factory=lambda: dict(_PRIVACY_FILTER_LABELS)
+    )
+
+
+class MLModelDetectorConfig(_PrivacyFilterDetectorConfig):
+    torch_dtype: str = "float16"
     batch_size: int = 1
-    label_mapping: dict[str, str] = Field(default_factory=lambda: {
-        "private_person": "PERSON",
-        "private_email": "EMAIL_ADDRESS",
-        "private_phone": "PHONE_NUMBER",
-        "private_address": "LOCATION",
-        "private_date": "DATE_TIME",
-        "private_url": "URL",
-        "account_number": "FINANCIAL",
-        "secret": "SECRET",
-    })
 
 
-class OnnxDetectorConfig(BaseModel):
-    enabled: bool = False
-    model_name: str = "openai/privacy-filter"
-    # Pin a Hugging Face revision (tag or commit sha); None = latest.
-    revision: str | None = None
-    # Off by default; only the tokenizer loads from the repo and the default
-    # model needs no repo code.
-    trust_remote_code: bool = False
+class OnnxDetectorConfig(_PrivacyFilterDetectorConfig):
     onnx_variant: str = "q4f16"
-    device: str = "auto"
-    score_threshold: float = 0.5
     max_length: int = 128000
     cache_dir: str | None = None
-    label_mapping: dict[str, str] = Field(default_factory=lambda: {
-        "private_person": "PERSON",
-        "private_email": "EMAIL_ADDRESS",
-        "private_phone": "PHONE_NUMBER",
-        "private_date": "DATE_TIME",
-        "private_address": "LOCATION",
-        "private_url": "URL",
-        "account_number": "FINANCIAL",
-        "secret": "SECRET",
-    })
 
 
 class BertNERDetectorConfig(BaseModel):
diff --git a/privaite/middleware/limits.py b/privaite/middleware/limits.py
index be8d89f..18e8d53 100644
--- a/privaite/middleware/limits.py
+++ b/privaite/middleware/limits.py
@@ -5,22 +5,8 @@
 from starlette.types import ASGIApp, Message, Receive, Scope, Send
 
 
-def _replay_body(body: bytes, receive: Receive) -> Receive:
-    """Yield the already-buffered body once, then delegate to the real receive."""
-    sent = False
-
-    async def _receive() -> Message:
-        nonlocal sent
-        if not sent:
-            sent = True
-            return {"type": "http.request", "body": body, "more_body": False}
-        return await receive()
-
-    return _receive
-
-
-def _replay_message(message: Message, receive: Receive) -> Receive:
-    """Yield one buffered control message once, then delegate to the real receive."""
+def _replay_once(message: Message, receive: Receive) -> Receive:
+    """Yield one already-buffered message once, then delegate to the real receive."""
     sent = False
 
     async def _receive() -> Message:
@@ -68,7 +54,7 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
             if message["type"] != "http.request":
                 # A control message such as http.disconnect: hand back control
                 # with the message replayed first so the app still sees it.
-                await self.app(scope, _replay_message(message, receive), send)
+                await self.app(scope, _replay_once(message, receive), send)
                 return
             chunk = message.get("body", b"")
             # Check the projected size before extending so a single oversized
@@ -79,7 +65,14 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
             body.extend(chunk)
             more_body = message.get("more_body", False)
 
-        await self.app(scope, _replay_body(bytes(body), receive), send)
+        await self.app(
+            scope,
+            _replay_once(
+                {"type": "http.request", "body": bytes(body), "more_body": False},
+                receive,
+            ),
+            send,
+        )
 
     async def _too_large(self, send: Send) -> None:
         payload = json.dumps(
diff --git a/privaite/pii/detector_base.py b/privaite/pii/detector_base.py
index 6b62e94..825a442 100644
--- a/privaite/pii/detector_base.py
+++ b/privaite/pii/detector_base.py
@@ -1,6 +1,9 @@
 from __future__ import annotations
 
+import asyncio
 from abc import ABC, abstractmethod
+from collections.abc import Callable
+from typing import Any
 
 from privaite.pii.entity import PIIEntity
 
@@ -48,3 +51,113 @@ def chunk_text(
             break
         start = max(end - overlap, start + 1)
     return chunks
+
+
+# One raw model result parsed to (label, score, chunk_start, chunk_end, span_text),
+# or None to skip it. Positions are relative to the chunk; the loop offsets them.
+NerResult = tuple[str, float, int, int, str]
+
+
+async def windowed_ner_detect(
+    text: str,
+    *,
+    predict: Callable[[str], list[Any]],
+    extract: Callable[[Any, str], NerResult | None],
+    label_mapping: dict[str, str],
+    score_threshold: float,
+    source: str,
+    max_chars: int = 1500,
+) -> list[PIIEntity]:
+    """Shared windowed detection loop for the ML NER detectors (BERT, GLiNER, the
+    HF pipeline). Each has the same shape: window the text so a model that truncates
+    long inputs still sees all of it, run the model per window, then threshold, map
+    the label, offset the span back to the full text, dedupe on (start, end, type),
+    and build a PIIEntity. ``predict(chunk)`` runs the model on one window (offloaded
+    to a thread); ``extract(result, chunk)`` turns one raw result into
+    (label, score, start, end, span_text) or None to skip it."""
+    entities: list[PIIEntity] = []
+    seen: set[tuple[int, int, str]] = set()
+    for offset, chunk in chunk_text(text, max_chars=max_chars):
+        results = await asyncio.to_thread(predict, chunk)
+        for result in results:
+            parsed = extract(result, chunk)
+            if parsed is None:
+                continue
+            label, score, start, end, span_text = parsed
+            if score < score_threshold:
+                continue
+            mapped_type = label_mapping.get(label)
+            if not mapped_type:
+                continue
+            start += offset
+            end += offset
+            key = (start, end, mapped_type)
+            if key in seen:
+                continue
+            seen.add(key)
+            entities.append(
+                PIIEntity(
+                    entity_type=mapped_type,
+                    text=span_text,
+                    start=start,
+                    end=end,
+                    score=float(score),
+                    source=source,
+                )
+            )
+    return entities
+
+
+def hf_pipeline_extract(result: dict, chunk: str) -> NerResult:
+    """Parse one Hugging Face token-classification result (aggregation_strategy
+    "simple"): entity_group/score/start/end, word stripped, falling back to the raw
+    chunk span. Shared by the BERT and generic ML pipeline detectors."""
+    start = result.get("start", 0)
+    end = result.get("end", 0)
+    word = result.get("word", chunk[start:end])
+    return (
+        result.get("entity_group", ""),
+        result.get("score", 0.0),
+        start,
+        end,
+        word.strip(),
+    )
+
+
+def resolve_torch_device(device_str: str) -> str:
+    """Map "auto" to the best available torch device (mps, then cuda, else cpu) and
+    pass any explicit device string through. Shared by the torch ``.to(device)``
+    detectors (GLiNER, the ML model); the HF-pipeline device format differs and
+    stays local to that detector."""
+    if device_str == "auto":
+        import torch
+
+        if torch.backends.mps.is_available():
+            return "mps"
+        if torch.cuda.is_available():
+            return "cuda"
+        return "cpu"
+    return device_str
+
+
+class HFPipelineDetector(PIIDetector):
+    """Shared ``detect()`` for the Hugging Face token-classification detectors (BERT
+    NER and the generic ML pipeline). Both build a ``self._classifier`` pipeline in
+    initialize() and differ only in how that pipeline is loaded, so detection itself
+    (window, run, threshold, map, offset, dedupe) lives here once. ``source`` is the
+    detector's own name."""
+
+    config: Any
+    _classifier: Any
+
+    async def detect(self, text: str, language: str = "en") -> list[PIIEntity]:
+        if self._classifier is None:
+            raise RuntimeError(f"{type(self).__name__} not initialized")
+        return await windowed_ner_detect(
+            text,
+            predict=self._classifier,
+            extract=hf_pipeline_extract,
+            label_mapping=self.config.label_mapping,
+            score_threshold=self.config.score_threshold,
+            source=self.name,
+        )
diff --git a/privaite/pii/detector_bert_ner.py b/privaite/pii/detector_bert_ner.py
index 044d778..c239438 100644
--- a/privaite/pii/detector_bert_ner.py
+++ b/privaite/pii/detector_bert_ner.py
@@ -5,13 +5,12 @@
 from typing import Any
 
 from privaite.config.schema import BertNERDetectorConfig
-from privaite.pii.detector_base import PIIDetector, chunk_text
-from privaite.pii.entity import PIIEntity
+from privaite.pii.detector_base import HFPipelineDetector, resolve_torch_device
 
 logger = logging.getLogger("privaite.pii.detector_bert_ner")
 
 
-class BertNERDetector(PIIDetector):
+class BertNERDetector(HFPipelineDetector):
     def __init__(self, config: BertNERDetectorConfig) -> None:
         self.config = config
         self._classifier: Any = None
@@ -37,67 +36,12 @@ def _load() -> None:
         await asyncio.to_thread(_load)
         logger.info("BERT NER model loaded")
 
-    async def detect(self, text: str, language: str = "en") -> list[PIIEntity]:
-        if self._classifier is None:
-            raise RuntimeError("BertNERDetector not initialized")
-
-        # The HF pipeline silently truncates at model_max_length (~512 tokens),
-        # making PII past that point invisible; run overlapping windows instead.
-        pii_entities: list[PIIEntity] = []
-        seen: set[tuple[int, int, str]] = set()
-        for offset, chunk in chunk_text(text):
-            results = await asyncio.to_thread(self._classifier, chunk)
-
-            for result in results:
-                entity_group = result.get("entity_group", "")
-                score = result.get("score", 0.0)
-
-                if score < self.config.score_threshold:
-                    continue
-
-                mapped_type = self.config.label_mapping.get(entity_group)
-                if not mapped_type:
-                    continue
-
-                start = result.get("start", 0) + offset
-                end = result.get("end", 0) + offset
-                word = result.get("word", text[start:end])
-
-                key = (start, end, mapped_type)
-                if key in seen:
-                    continue
-                seen.add(key)
-
-                pii_entities.append(
-                    PIIEntity(
-                        entity_type=mapped_type,
-                        text=word.strip(),
-                        start=start,
-                        end=end,
-                        score=score,
-                        source="bert_ner",
-                    )
-                )
-
-        return pii_entities
-
     async def shutdown(self) -> None:
         self._classifier = None
 
 
 def _resolve_device(device_str: str) -> int | str:
-    if device_str == "auto":
-        import torch
-
-        if torch.backends.mps.is_available():
-            return "mps"
-        if torch.cuda.is_available():
-            return 0
-        return -1
-    if device_str == "cpu":
-        return -1
-    if device_str == "mps":
-        return "mps"
-    if device_str == "cuda":
-        return 0
-    return -1
+    # HF pipelines want a device index (-1 cpu, 0 cuda) or the "mps" string; map the
+    # shared torch device name onto that format.
+    pipeline_device: dict[str, int | str] = {"cpu": -1, "cuda": 0, "mps": "mps"}
+    return pipeline_device.get(resolve_torch_device(device_str), -1)
diff --git a/privaite/pii/detector_gliner.py b/privaite/pii/detector_gliner.py
index 402b1e2..373ea57 100644
--- a/privaite/pii/detector_gliner.py
+++ b/privaite/pii/detector_gliner.py
@@ -5,12 +5,23 @@
 from typing import Any
 
 from privaite.config.schema import GlinerDetectorConfig
-from privaite.pii.detector_base import PIIDetector, chunk_text
+from privaite.pii.detector_base import (
+    NerResult,
+    PIIDetector,
+    resolve_torch_device,
+    windowed_ner_detect,
+)
 from privaite.pii.entity import PIIEntity
 
 logger = logging.getLogger("privaite.pii.detector_gliner")
 
 
+def _gliner_extract(result: dict, chunk: str) -> NerResult:
+    start = int(result.get("start", 0))
+    end = int(result.get("end", 0))
+    return (result.get("label", ""), result.get("score", 0.0), start, end, chunk[start:end])
+
+
 class GlinerDetector(PIIDetector):
     """GLiNER PII detector (the `max` preset).
 
@@ -42,7 +53,7 @@ def _load() -> None:
             model = GLiNER.from_pretrained(
                 self.config.model_name, revision=self.config.revision
             )
-            device = _resolve_device(self.config.device)
+            device = resolve_torch_device(self.config.device)
             self._model = model.to(device)
 
         logger.info("Loading GLiNER model: %s ...", self.config.model_name)
@@ -53,58 +64,22 @@ async def detect(self, text: str, language: str = "en") -> list[PIIEntity]:
         if self._model is None:
             raise RuntimeError("GlinerDetector not initialized")
 
+        # GLiNER truncates long inputs (~384 tokens); windowing (max_chars=1200)
+        # keeps PII past that point visible. It filters by threshold itself, and the
+        # shared loop re-checks it and dedupes overlapping-window hits.
         labels = self.config.labels
         threshold = self.config.score_threshold
-        pii_entities: list[PIIEntity] = []
-        seen: set[tuple[int, int, str]] = set()
-
-        # GLiNER truncates long inputs (~384 tokens); run overlapping windows so PII
-        # past that point is still seen, and offset spans back to the full text.
-        for offset, chunk in chunk_text(text, max_chars=1200):
-            results = await asyncio.to_thread(
-                self._model.predict_entities, chunk, labels, threshold=threshold
-            )
-            for result in results:
-                score = result.get("score", 0.0)
-                if score < threshold:
-                    continue
-
-                mapped_type = self.config.label_mapping.get(result.get("label", ""))
-                if not mapped_type:
-                    continue
-
-                start = int(result.get("start", 0)) + offset
-                end = int(result.get("end", 0)) + offset
-
-                key = (start, end, mapped_type)
-                if key in seen:
-                    continue
-                seen.add(key)
-
-                pii_entities.append(
-                    PIIEntity(
-                        entity_type=mapped_type,
-                        text=text[start:end],
-                        start=start,
-                        end=end,
-                        score=float(score),
-                        source="gliner",
-                    )
-                )
-
-        return pii_entities
+        return await windowed_ner_detect(
+            text,
+            predict=lambda chunk: self._model.predict_entities(
+                chunk, labels, threshold=threshold
+            ),
+            extract=_gliner_extract,
+            label_mapping=self.config.label_mapping,
+            score_threshold=threshold,
+            source="gliner",
+            max_chars=1200,
+        )
 
     async def shutdown(self) -> None:
         self._model = None
-
-
-def _resolve_device(device_str: str) -> str:
-    if device_str == "auto":
-        import torch
-
-        if torch.backends.mps.is_available():
-            return "mps"
-        if torch.cuda.is_available():
-            return "cuda"
-        return "cpu"
-    return device_str
diff --git a/privaite/pii/detector_mlmodel.py b/privaite/pii/detector_mlmodel.py
index 8480f6a..141da31 100644
--- a/privaite/pii/detector_mlmodel.py
+++ b/privaite/pii/detector_mlmodel.py
@@ -5,24 +5,11 @@
 from typing import Any
 
 from privaite.config.schema import MLModelDetectorConfig
-from privaite.pii.detector_base import PIIDetector, chunk_text
-from privaite.pii.entity import PIIEntity
+from privaite.pii.detector_base import HFPipelineDetector, resolve_torch_device
 
 logger = logging.getLogger("privaite.pii.detector_mlmodel")
 
 
-def _resolve_device(device_str: str) -> str:
-    if device_str == "auto":
-        import torch
-
-        if torch.backends.mps.is_available():
-            return "mps"
-        if torch.cuda.is_available():
-            return "cuda"
-        return "cpu"
-    return device_str
-
-
 def _resolve_dtype(dtype_str: str):
     import torch
 
@@ -31,7 +18,7 @@ def _resolve_dtype(dtype_str: str):
     )
 
 
-class MLModelDetector(PIIDetector):
+class MLModelDetector(HFPipelineDetector):
     def __init__(self, config: MLModelDetectorConfig) -> None:
         self.config = config
         self._classifier: Any = None
@@ -44,7 +31,7 @@ async def initialize(self) -> None:
         def _load() -> None:
             from transformers import AutoModelForTokenClassification, AutoTokenizer, pipeline
 
-            device = _resolve_device(self.config.device)
+            device = resolve_torch_device(self.config.device)
             dtype = _resolve_dtype(self.config.torch_dtype)
 
             logger.info(
@@ -80,49 +67,5 @@ def _load() -> None:
         await asyncio.to_thread(_load)
         logger.info("ML model loaded successfully")
 
-    async def detect(self, text: str, language: str = "en") -> list[PIIEntity]:
-        if self._classifier is None:
-            raise RuntimeError("MLModelDetector not initialized")
-
-        # The HF pipeline silently truncates at model_max_length, making PII past
-        # that point invisible; run overlapping windows instead.
-        pii_entities: list[PIIEntity] = []
-        seen: set[tuple[int, int, str]] = set()
-        for offset, chunk in chunk_text(text):
-            results = await asyncio.to_thread(self._classifier, chunk)
-
-            for result in results:
-                entity_group = result.get("entity_group", "")
-                score = result.get("score", 0.0)
-
-                if score < self.config.score_threshold:
-                    continue
-
-                mapped_type = self.config.label_mapping.get(entity_group)
-                if not mapped_type:
-                    continue
-
-                start = result.get("start", 0) + offset
-                end = result.get("end", 0) + offset
-                word = result.get("word", text[start:end])
-
-                key = (start, end, mapped_type)
-                if key in seen:
-                    continue
-                seen.add(key)
-
-                pii_entities.append(
-                    PIIEntity(
-                        entity_type=mapped_type,
-                        text=word.strip(),
-                        start=start,
-                        end=end,
-                        score=score,
-                        source="mlmodel",
-                    )
-                )
-
-        return pii_entities
-
     async def shutdown(self) -> None:
         self._classifier = None
diff --git a/privaite/pii/recognizer_base.py b/privaite/pii/recognizer_base.py
new file mode 100644
index 0000000..f627841
--- /dev/null
+++ b/privaite/pii/recognizer_base.py
@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+import re
+
+from presidio_analyzer import EntityRecognizer, RecognizerResult
+from presidio_analyzer.nlp_engine import NlpArtifacts
+
+
+class RegexSpanRecognizer(EntityRecognizer):
+    """Base for recognizers that emit one RecognizerResult per regex match over a
+    fixed list of patterns. Subclasses fill ``self._specs`` with
+    (compiled, entity_type, score, group) tuples in __init__; ``group`` is the named
+    span to report, or None for the whole match. Recognizers with per-match logic
+    (trimming, normalization) implement their own ``analyze`` instead."""
+
+    _specs: list[tuple[re.Pattern[str], str, float, str | None]]
+
+    def load(self) -> None:
+        pass
+
+    def analyze(
+        self, text: str, entities: list[str], nlp_artifacts: NlpArtifacts | None = None
+    ) -> list[RecognizerResult]:
+        results = []
+        for compiled, entity_type, score, group in self._specs:
+            for match in compiled.finditer(text):
+                results.append(
+                    RecognizerResult(
+                        entity_type=entity_type,
+                        start=match.start(group) if group else match.start(),
+                        end=match.end(group) if group else match.end(),
+                        score=score,
+                    )
+                )
+        return results
diff --git a/privaite/pii/recognizer_custom.py b/privaite/pii/recognizer_custom.py
index 9ffbe41..1730953 100644
--- a/privaite/pii/recognizer_custom.py
+++ b/privaite/pii/recognizer_custom.py
@@ -2,44 +2,22 @@
 
 import re
 
-from presidio_analyzer import EntityRecognizer, RecognizerResult
-from presidio_analyzer.nlp_engine import NlpArtifacts
-
 from privaite.config.schema import CustomPatternConfig
+from privaite.pii.recognizer_base import RegexSpanRecognizer
 
 
-class CustomPatternRecognizer(EntityRecognizer):
+class CustomPatternRecognizer(RegexSpanRecognizer):
     def __init__(
         self,
         patterns: list[CustomPatternConfig],
         supported_language: str = "fr",
     ) -> None:
-        entity_types = list({p.entity_type for p in patterns})
         super().__init__(
-            supported_entities=entity_types,
+            supported_entities=list({p.entity_type for p in patterns}),
             supported_language=supported_language,
             name="CustomPatternRecognizer",
         )
-        self._patterns = [
-            (re.compile(p.pattern, re.IGNORECASE | re.UNICODE), p.entity_type, p.score)
+        self._specs = [
+            (re.compile(p.pattern, re.IGNORECASE | re.UNICODE), p.entity_type, p.score, None)
             for p in patterns
         ]
-
-    def load(self) -> None:
-        pass
-
-    def analyze(
-        self, text: str, entities: list[str], nlp_artifacts: NlpArtifacts | None = None
-    ) -> list[RecognizerResult]:
-        results = []
-        for compiled, entity_type, score in self._patterns:
-            for match in compiled.finditer(text):
-                results.append(
-                    RecognizerResult(
-                        entity_type=entity_type,
-                        start=match.start(),
-                        end=match.end(),
-                        score=score,
-                    )
-                )
-        return results
diff --git a/privaite/pii/recognizer_fr_date.py b/privaite/pii/recognizer_fr_date.py
index 60e79a4..ff23956 100644
--- a/privaite/pii/recognizer_fr_date.py
+++ b/privaite/pii/recognizer_fr_date.py
@@ -2,8 +2,7 @@
 
 import re
 
-from presidio_analyzer import EntityRecognizer, RecognizerResult
-from presidio_analyzer.nlp_engine import NlpArtifacts
+from privaite.pii.recognizer_base import RegexSpanRecognizer
 
 MONTHS_FR = (
     "janvier|février|fevrier|mars|avril|mai|juin|"
@@ -29,34 +28,11 @@
 COMPILED = [re.compile(p, re.IGNORECASE | re.UNICODE) for p in PATTERNS]
 
 
-class FrenchDateRecognizer(EntityRecognizer):
+class FrenchDateRecognizer(RegexSpanRecognizer):
     def __init__(self, supported_language: str = "fr") -> None:
         super().__init__(
             supported_entities=["DATE_TIME"],
             supported_language=supported_language,
             name="FrenchDateRecognizer",
         )
-
-    def load(self) -> None:
-        pass
-
-    def analyze(
-        self, text: str, entities: list[str], nlp_artifacts: NlpArtifacts | None = None
-    ) -> list[RecognizerResult]:
-        results = []
-
-        for pattern in COMPILED:
-            for match in pattern.finditer(text):
-                start = match.start("date")
-                end = match.end("date")
-
-                results.append(
-                    RecognizerResult(
-                        entity_type="DATE_TIME",
-                        start=start,
-                        end=end,
-                        score=0.85,
-                    )
-                )
-
-        return results
+        self._specs = [(c, "DATE_TIME", 0.85, "date") for c in COMPILED]
diff --git a/privaite/pii/recognizer_location.py b/privaite/pii/recognizer_location.py
index 0b4d374..7f206b6 100644
--- a/privaite/pii/recognizer_location.py
+++ b/privaite/pii/recognizer_location.py
@@ -2,8 +2,7 @@
 
 import re
 
-from presidio_analyzer import EntityRecognizer, RecognizerResult
-from presidio_analyzer.nlp_engine import NlpArtifacts
+from privaite.pii.recognizer_base import RegexSpanRecognizer
 
 _LOC_GROUP = r"(?P<loc>[A-ZÀ-Ÿ][A-ZÀ-Ÿa-zà-ÿ\-]+(?:[\s\-]+[A-ZÀ-Ÿa-zà-ÿ\-]+){0,3})"
 
@@ -28,31 +27,11 @@
 ]
 
 
-class ContextualLocationRecognizer(EntityRecognizer):
+class ContextualLocationRecognizer(RegexSpanRecognizer):
     def __init__(self, supported_language: str = "fr") -> None:
         super().__init__(
             supported_entities=["LOCATION"],
             supported_language=supported_language,
             name="ContextualLocationRecognizer",
         )
-
-    def load(self) -> None:
-        pass
-
-    def analyze(
-        self, text: str, entities: list[str], nlp_artifacts: NlpArtifacts | None = None
-    ) -> list[RecognizerResult]:
-        results = []
-        for pattern in COMPILED:
-            for match in pattern.finditer(text):
-                start = match.start("loc")
-                end = match.end("loc")
-                results.append(
-                    RecognizerResult(
-                        entity_type="LOCATION",
-                        start=start,
-                        end=end,
-                        score=0.95,
-                    )
-                )
-        return results
+        self._specs = [(c, "LOCATION", 0.95, "loc") for c in COMPILED]
diff --git a/tests/test_api/test_limits.py b/tests/test_api/test_limits.py
index 4356463..58f5390 100644
--- a/tests/test_api/test_limits.py
+++ b/tests/test_api/test_limits.py
@@ -149,3 +149,71 @@ async def send(message):
     await mw({"type": "http", "headers": []}, receive, send)
 
     assert bytes(received) == b"hello world"
+
+
+@pytest.mark.asyncio
+async def test_non_http_scope_passes_through():
+    """A non-http scope (websocket, lifespan) bypasses the size check untouched."""
+    seen = {}
+
+    async def downstream(scope, receive, send):
+        seen["type"] = scope["type"]
+
+    mw = RequestSizeLimitMiddleware(downstream, max_bytes=10)
+
+    await mw({"type": "lifespan"}, None, None)
+
+    assert seen["type"] == "lifespan"
+
+
+@pytest.mark.asyncio
+async def test_non_integer_content_length_is_ignored():
+    """A malformed Content-Length header does not crash; the body is still counted."""
+    received = bytearray()
+
+    async def downstream(scope, receive, send):
+        while True:
+            message = await receive()
+            received.extend(message.get("body", b""))
+            if not message.get("more_body", False):
+                break
+
+    mw = RequestSizeLimitMiddleware(downstream, max_bytes=1000)
+    chunks = [{"type": "http.request", "body": b"hi", "more_body": False}]
+
+    async def receive():
+        return chunks.pop(0)
+
+    async def send(message):
+        pass
+
+    await mw(
+        {"type": "http", "headers": [(b"content-length", b"not-a-number")]},
+        receive,
+        send,
+    )
+
+    assert bytes(received) == b"hi"
+
+
+@pytest.mark.asyncio
+async def test_control_message_is_replayed_to_app():
+    """An http.disconnect arriving before the body is replayed to the downstream app."""
+    seen = []
+
+    async def downstream(scope, receive, send):
+        message = await receive()
+        seen.append(message["type"])
+
+    mw = RequestSizeLimitMiddleware(downstream, max_bytes=1000)
+    chunks = [{"type": "http.disconnect"}]
+
+    async def receive():
+        return chunks.pop(0)
+
+    async def send(message):
+        pass
+
+    await mw({"type": "http", "headers": []}, receive, send)
+
+    assert seen == ["http.disconnect"]
diff --git a/tests/test_pii/test_windowed_detect.py b/tests/test_pii/test_windowed_detect.py
new file mode 100644
index 0000000..aa4fbf9
--- /dev/null
+++ b/tests/test_pii/test_windowed_detect.py
@@ -0,0 +1,150 @@
+from __future__ import annotations
+
+import sys
+import types
+
+import pytest
+
+from privaite.pii.detector_base import (
+    hf_pipeline_extract,
+    resolve_torch_device,
+    windowed_ner_detect,
+)
+
+
+def _extract(result, chunk):
+    """Minimal extract: result is a dict with label/score/start/end (chunk-relative)."""
+    if result.get("skip"):
+        return None
+    start, end = result["start"], result["end"]
+    return (result["label"], result["score"], start, end, chunk[start:end])
+
+
+@pytest.mark.asyncio
+async def test_maps_label_and_builds_entity():
+    def predict(chunk):
+        return [{"label": "M", "score": 0.9, "start": 0, "end": 4}]
+
+    out = await windowed_ner_detect(
+        "MARK rest",
+        predict=predict,
+        extract=_extract,
+        label_mapping={"M": "PERSON"},
+        score_threshold=0.5,
+        source="test",
+    )
+    assert len(out) == 1
+    e = out[0]
+    assert (e.entity_type, e.text, e.start, e.end, e.source) == (
+        "PERSON", "MARK", 0, 4, "test",
+    )
+    assert e.score == 0.9
+
+
+@pytest.mark.asyncio
+async def test_below_threshold_unmapped_and_none_are_skipped():
+    def predict(chunk):
+        return [
+            {"label": "M", "score": 0.2, "start": 0, "end": 4},   # below threshold
+            {"label": "UNKNOWN", "score": 0.9, "start": 0, "end": 4},  # unmapped
+            {"skip": True},                                        # extract -> None
+        ]
+
+    out = await windowed_ner_detect(
+        "MARK",
+        predict=predict,
+        extract=_extract,
+        label_mapping={"M": "PERSON"},
+        score_threshold=0.5,
+        source="test",
+    )
+    assert out == []
+
+
+@pytest.mark.asyncio
+async def test_same_span_in_overlapping_windows_is_deduped():
+    # Long enough to window (max_chars < len), no spaces so chunks are fixed-size;
+    # "MARK" sits in the overlap of the first two 300-char windows (overlap 200),
+    # so it is found twice and offset back to the same (150, 154) span -> one entity.
+    text = "x" * 150 + "MARK" + "y" * 350
+
+    def predict(chunk):
+        i = chunk.find("MARK")
+        return [{"label": "M", "score": 0.9, "start": i, "end": i + 4}] if i != -1 else []
+
+    out = await windowed_ner_detect(
+        text,
+        predict=predict,
+        extract=_extract,
+        label_mapping={"M": "PERSON"},
+        score_threshold=0.5,
+        source="test",
+        max_chars=300,
+    )
+    assert len(out) == 1
+    assert (out[0].start, out[0].end, out[0].text) == (150, 154, "MARK")
+
+
+def test_hf_pipeline_extract_uses_word_stripped_with_span_fallback():
+    # word present: stripped and returned as span text.
+    label, score, start, end, span = hf_pipeline_extract(
+        {"entity_group": "PER", "score": 0.8, "start": 2, "end": 6, "word": " Bob "},
+        "hi Bob!",
+    )
+    assert (label, score, start, end, span) == ("PER", 0.8, 2, 6, "Bob")
+
+    # word absent: falls back to the raw chunk span.
+    _, _, _, _, span2 = hf_pipeline_extract(
+        {"entity_group": "PER", "score": 0.8, "start": 3, "end": 6}, "hi Bob!"
+    )
+    assert span2 == "Bob"
+
+
+@pytest.mark.asyncio
+async def test_mlmodel_detector_maps_and_labels_its_source():
+    from privaite.config.schema import MLModelDetectorConfig
+    from privaite.pii.detector_mlmodel import MLModelDetector
+
+    detector = MLModelDetector(
+        MLModelDetectorConfig(label_mapping={"PER": "PERSON"}, score_threshold=0.5)
+    )
+    detector._classifier = lambda chunk: [
+        {"entity_group": "PER", "score": 0.9, "start": 3, "end": 6, "word": "Bob"}
+    ]
+
+    out = await detector.detect("hi Bob")
+
+    assert len(out) == 1
+    assert out[0].entity_type == "PERSON"
+    assert out[0].source == "mlmodel"
+
+
+@pytest.mark.asyncio
+async def test_mlmodel_detector_requires_initialization():
+    from privaite.config.schema import MLModelDetectorConfig
+    from privaite.pii.detector_mlmodel import MLModelDetector
+
+    with pytest.raises(RuntimeError, match="not initialized"):
+        await MLModelDetector(MLModelDetectorConfig()).detect("text")
+
+
+def _fake_torch(mps: bool, cuda: bool) -> types.ModuleType:
+    torch = types.ModuleType("torch")
+    torch.backends = types.SimpleNamespace(  # type: ignore[attr-defined]
+        mps=types.SimpleNamespace(is_available=lambda: mps)
+    )
+    torch.cuda = types.SimpleNamespace(is_available=lambda: cuda)  # type: ignore[attr-defined]
+    return torch
+
+
+@pytest.mark.parametrize(
+    "mps,cuda,expected",
+    [(True, False, "mps"), (False, True, "cuda"), (False, False, "cpu")],
+)
+def test_resolve_torch_device_auto(monkeypatch, mps, cuda, expected):
+    monkeypatch.setitem(sys.modules, "torch", _fake_torch(mps, cuda))
+    assert resolve_torch_device("auto") == expected
+
+
+def test_resolve_torch_device_passes_explicit_through():
+    assert resolve_torch_device("cpu") == "cpu"

From d032f68ac942782e4b296c0421ca76f6e455102e Mon Sep 17 00:00:00 2001
From: crp4222 <121295348+crp4222@users.noreply.github.com>
Date: Thu, 9 Jul 2026 15:33:13 +0100
Subject: [PATCH 2/4] refactor: extract streaming drain, endpoint helpers and
 fuzzy matcher

---
 privaite/api/chat.py          |  47 ++++++--------
 privaite/api/completions.py   |  63 ++++++++----------
 privaite/api/pipeline.py      |  37 +++++++++++
 privaite/pii/deanonymizer.py  |  66 +++++++++++--------
 privaite/streaming/handler.py | 116 ++++++++++++++++++----------------
 5 files changed, 183 insertions(+), 146 deletions(-)

diff --git a/privaite/api/chat.py b/privaite/api/chat.py
index cc58aff..73eb100 100644
--- a/privaite/api/chat.py
+++ b/privaite/api/chat.py
@@ -12,10 +12,10 @@
     record_pii_stats,
 )
 from privaite.api.pipeline import (
-    anonymize_or_error,
     call_provider,
     dump_response,
-    sse_response,
+    resolve_input,
+    stream_or_error,
     validate_model,
 )
 from privaite.config.schema import PrivAiTeConfig
@@ -66,43 +66,32 @@ async def chat_completions(
     if error is not None:
         return error
 
-    mapping = None
+    async def _anonymize() -> tuple[list, Any]:
+        anon, mapping = await pii_engine.process_request(messages)
+        record_pii_stats(request, mapping)
+        return anon, mapping
 
-    if config.pii.enabled and pii_engine is not None:
-
-        async def _anonymize() -> tuple[list, Any]:
-            anon, mapping = await pii_engine.process_request(messages)
-            record_pii_stats(request, mapping)
-            return anon, mapping
-
-        result, error = await anonymize_or_error(_anonymize, config, logger)
-        if error is not None:
-            return error
-        if result is not None:
-            messages, mapping = result
+    messages, mapping, error = await resolve_input(
+        config.pii.enabled and pii_engine is not None, messages, _anonymize, config, logger
+    )
+    if error is not None:
+        return error
 
     kwargs = {k: v for k, v in body.items() if k not in ("model", "messages", "stream")}
 
     if stream:
-        litellm_stream, error = await call_provider(
+        from privaite.streaming.handler import StreamingHandler
+
+        deanon_config = config.pii.deanonymization if config.pii.enabled else None
+        return await stream_or_error(
             lambda: provider_router.streaming_completion(
                 model_alias=model, messages=messages, **kwargs
             ),
+            lambda s: StreamingHandler.stream_response(
+                litellm_stream=s, mapping=mapping, deanonymizer_config=deanon_config
+            ),
             model, logger,
         )
-        if error is not None:
-            return error
-
-        from privaite.streaming.handler import StreamingHandler
-
-        deanon_config = config.pii.deanonymization if config.pii.enabled else None
-        return sse_response(
-            StreamingHandler.stream_response(
-                litellm_stream=litellm_stream,
-                mapping=mapping,
-                deanonymizer_config=deanon_config,
-            )
-        )
 
     response, error = await call_provider(
         lambda: provider_router.completion(model_alias=model, messages=messages, **kwargs),
diff --git a/privaite/api/completions.py b/privaite/api/completions.py
index 2057608..1fb5157 100644
--- a/privaite/api/completions.py
+++ b/privaite/api/completions.py
@@ -12,10 +12,10 @@
     record_pii_stats,
 )
 from privaite.api.pipeline import (
-    anonymize_or_error,
     call_provider,
     dump_response,
-    sse_response,
+    resolve_input,
+    stream_or_error,
     validate_model,
 )
 from privaite.config.schema import PrivAiTeConfig
@@ -42,50 +42,39 @@ async def completions(
     if error is not None:
         return error
 
-    mapping = None
-
-    if config.pii.enabled and pii_engine is not None:
-
-        async def _anonymize() -> tuple[Any, Any]:
-            if isinstance(prompt, list) and all(isinstance(p, str) for p in prompt):
-                msgs = [{"role": "user", "content": p} for p in prompt]
-                msgs, mapping = await pii_engine.process_request(msgs)
-                anon_prompt: Any = [m["content"] for m in msgs]
-            else:
-                msgs = [{"role": "user", "content": prompt}]
-                msgs, mapping = await pii_engine.process_request(msgs)
-                anon_prompt = msgs[0]["content"]
-            record_pii_stats(request, mapping)
-            return anon_prompt, mapping
-
-        result, error = await anonymize_or_error(_anonymize, config, logger)
-        if error is not None:
-            return error
-        if result is not None:
-            prompt, mapping = result
+    async def _anonymize() -> tuple[Any, Any]:
+        if isinstance(prompt, list) and all(isinstance(p, str) for p in prompt):
+            msgs = [{"role": "user", "content": p} for p in prompt]
+            msgs, mapping = await pii_engine.process_request(msgs)
+            anon_prompt: Any = [m["content"] for m in msgs]
+        else:
+            msgs = [{"role": "user", "content": prompt}]
+            msgs, mapping = await pii_engine.process_request(msgs)
+            anon_prompt = msgs[0]["content"]
+        record_pii_stats(request, mapping)
+        return anon_prompt, mapping
+
+    prompt, mapping, error = await resolve_input(
+        config.pii.enabled and pii_engine is not None, prompt, _anonymize, config, logger
+    )
+    if error is not None:
+        return error
 
     kwargs = {k: v for k, v in body.items() if k not in ("model", "prompt", "stream")}
 
     if stream:
-        litellm_stream, error = await call_provider(
+        from privaite.streaming.handler import StreamingHandler
+
+        deanon_config = config.pii.deanonymization if config.pii.enabled else None
+        return await stream_or_error(
             lambda: provider_router.streaming_text_completion(
                 model_alias=model, prompt=prompt, **kwargs
             ),
+            lambda s: StreamingHandler.stream_text_response(
+                litellm_stream=s, mapping=mapping, deanonymizer_config=deanon_config
+            ),
             model, logger,
         )
-        if error is not None:
-            return error
-
-        from privaite.streaming.handler import StreamingHandler
-
-        deanon_config = config.pii.deanonymization if config.pii.enabled else None
-        return sse_response(
-            StreamingHandler.stream_text_response(
-                litellm_stream=litellm_stream,
-                mapping=mapping,
-                deanonymizer_config=deanon_config,
-            )
-        )
 
     response, error = await call_provider(
         lambda: provider_router.text_completion(model_alias=model, prompt=prompt, **kwargs),
diff --git a/privaite/api/pipeline.py b/privaite/api/pipeline.py
index 2c33bab..978a2b4 100644
--- a/privaite/api/pipeline.py
+++ b/privaite/api/pipeline.py
@@ -83,12 +83,49 @@ async def call_provider(
         return None, provider_error_response(exc)
 
 
+async def resolve_input(
+    pii_enabled: bool,
+    raw: T,
+    anonymize: Callable[[], Awaitable[tuple[T, Any]]],
+    config: PrivAiTeConfig,
+    log: logging.Logger,
+) -> tuple[T, Any, JSONResponse | None]:
+    """Anonymize an endpoint's input under the fail-closed policy, returning
+    (input_to_forward, mapping, error). The input is the ``raw`` payload unchanged
+    when PII is off, when ``on_error: allow`` opts out, or on error (unused then);
+    otherwise it is the anonymized payload with its reversible mapping.
+    """
+    if not pii_enabled:
+        return raw, None, None
+    result, error = await anonymize_or_error(anonymize, config, log)
+    if error is not None:
+        return raw, None, error
+    if result is None:
+        return raw, None, None
+    payload, mapping = result
+    return payload, mapping, None
+
+
 def sse_response(generator: AsyncIterator[str]) -> StreamingResponse:
     return StreamingResponse(
         generator, media_type="text/event-stream", headers=dict(SSE_HEADERS)
     )
 
 
+async def stream_or_error(
+    provider_call: Callable[[], Awaitable[Any]],
+    make_stream: Callable[[Any], AsyncIterator[str]],
+    model: str,
+    log: logging.Logger,
+) -> StreamingResponse | JSONResponse:
+    """Open a provider stream (mapping any failure to the OpenAI error shape) and
+    wrap the restored generator as an SSE response."""
+    litellm_stream, error = await call_provider(provider_call, model, log)
+    if error is not None:
+        return error
+    return sse_response(make_stream(litellm_stream))
+
+
 def dump_response(response: Any) -> dict:
     if hasattr(response, "model_dump"):
         return response.model_dump()
diff --git a/privaite/pii/deanonymizer.py b/privaite/pii/deanonymizer.py
index 2b73b8a..93794c5 100644
--- a/privaite/pii/deanonymizer.py
+++ b/privaite/pii/deanonymizer.py
@@ -71,35 +71,47 @@ def _overlaps(start: int, end: int) -> bool:
                 continue
 
             for i in range(len(tokens) - window + 1):
-                raw_start = tokens[i][0]
-                raw_end = tokens[i + window - 1][1]
-
-                # Trim sentence punctuation off the edges: match and replace only
-                # the core, so "Michel Deu," keeps its comma and "<PERSON_3>." is
-                # still recognized as a placeholder shape.
-                start, end = raw_start, raw_end
-                while start < end and text[start] in _EDGE_PUNCT:
-                    start += 1
-                while end > start and text[end - 1] in _EDGE_PUNCT:
-                    end -= 1
-                if start >= end or _overlaps(start, end):
-                    continue
-
-                candidate = text[start:end]
-                if candidate == original or candidate in known_fakes:
-                    continue
-                if _PLACEHOLDER_SHAPE.match(candidate):
-                    continue
-                # Compare with normalized whitespace: a fake re-typed across a
-                # newline or a double space is exactly what fuzzy exists to catch,
-                # and the raw slice would otherwise pay a ratio penalty per
-                # whitespace character.
-                normalized = " ".join(candidate.split())
-                ratio = SequenceMatcher(None, normalized.lower(), fake.lower()).ratio()
-                if ratio >= self.config.fuzzy_threshold:
-                    replacements.append((start, end, original))
+                span = self._fuzzy_edge_match(
+                    text, tokens[i][0], tokens[i + window - 1][1], fake, original, known_fakes
+                )
+                if span is not None and not _overlaps(*span):
+                    replacements.append((span[0], span[1], original))
 
         for start, end, original in sorted(replacements, reverse=True):
             text = text[:start] + original + text[end:]
 
         return text
+
+    def _fuzzy_edge_match(
+        self,
+        text: str,
+        raw_start: int,
+        raw_end: int,
+        fake: str,
+        original: str,
+        known_fakes: set[str],
+    ) -> tuple[int, int] | None:
+        """Trim sentence punctuation off a token window and, if the core slice
+        fuzzily matches ``fake`` over the threshold, return its (start, end); else
+        None. The core is skipped when it is the real value, an exact fake, or a
+        placeholder shape ("Michel Deu," keeps its comma; "<PERSON_3>." stays a
+        placeholder). Whitespace is normalized before scoring so a fake re-typed
+        across a newline or double space still matches instead of paying a ratio
+        penalty per whitespace char."""
+        start, end = raw_start, raw_end
+        while start < end and text[start] in _EDGE_PUNCT:
+            start += 1
+        while end > start and text[end - 1] in _EDGE_PUNCT:
+            end -= 1
+        if start >= end:
+            return None
+        candidate = text[start:end]
+        if candidate == original or candidate in known_fakes:
+            return None
+        if _PLACEHOLDER_SHAPE.match(candidate):
+            return None
+        normalized = " ".join(candidate.split())
+        ratio = SequenceMatcher(None, normalized.lower(), fake.lower()).ratio()
+        if ratio >= self.config.fuzzy_threshold:
+            return (start, end)
+        return None
diff --git a/privaite/streaming/handler.py b/privaite/streaming/handler.py
index 0f6a1aa..4030720 100644
--- a/privaite/streaming/handler.py
+++ b/privaite/streaming/handler.py
@@ -2,7 +2,7 @@
 
 import json
 import logging
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterator, Callable, Iterator
 from typing import Any
 
 from privaite.config.schema import DeanonymizationConfig
@@ -21,6 +21,24 @@
 _REASONING_FIELDS = ("reasoning_content", "reasoning")
 
 
+def _drain(
+    buffers: dict[Any, Any],
+    finished: set[int],
+    make_chunk: Callable[[Any, str], dict],
+) -> Iterator[str]:
+    """After the stream ends, emit whatever each buffer still holds so no restored
+    text is silently dropped, skipping choices that already got a finish chunk. The
+    buffer key is the choice index, or a tuple whose first element is it;
+    ``make_chunk(key, remaining)`` builds the chunk dict to send."""
+    for key, buf in buffers.items():
+        idx = key[0] if isinstance(key, tuple) else key
+        if idx in finished:
+            continue
+        remaining = buf.flush()
+        if remaining:
+            yield format_sse_event(json.dumps(make_chunk(key, remaining)))
+
+
 class StreamingHandler:
     @staticmethod
     async def stream_response(
@@ -197,44 +215,38 @@ def _flush_into_delta(idx: int, delta: dict, pre_events: list[dict]) -> None:
                 if visible:
                     yield format_sse_event(json.dumps(chunk_dict))
 
-            # Stream ended without a finish chunk for some choice: emit whatever
-            # the buffers still hold so no restored text is silently dropped.
-            for idx, buf in content_bufs.items():
-                if idx in finished:
-                    continue
-                remaining = buf.flush()
-                if remaining:
-                    yield format_sse_event(json.dumps(create_chunk_dict(
-                        content=remaining, model=model_name, index=idx
-                    )))
-            for (idx, field), buf in reasoning_bufs.items():
-                if idx in finished:
-                    continue
-                remaining = buf.flush()
-                if remaining:
-                    yield format_sse_event(json.dumps(create_delta_chunk(
-                        {field: remaining}, model=model_name, index=idx
-                    )))
-            for (idx, t_idx), buf in tool_bufs.items():
-                if idx in finished:
-                    continue
-                remaining = buf.flush()
-                if remaining:
-                    yield format_sse_event(json.dumps(create_delta_chunk(
-                        {"tool_calls": [
-                            {"index": t_idx, "function": {"arguments": remaining}}
-                        ]},
-                        model=model_name, index=idx,
-                    )))
-            for idx, buf in func_bufs.items():
-                if idx in finished:
-                    continue
-                remaining = buf.flush()
-                if remaining:
-                    yield format_sse_event(json.dumps(create_delta_chunk(
-                        {"function_call": {"arguments": remaining}},
-                        model=model_name, index=idx,
-                    )))
+            # Stream ended without a finish chunk for some choice: emit whatever the
+            # buffers still hold so no restored text is silently dropped.
+            for sse in _drain(
+                content_bufs,
+                finished,
+                lambda idx, r: create_chunk_dict(content=r, model=model_name, index=idx),
+            ):
+                yield sse
+            for sse in _drain(
+                reasoning_bufs,
+                finished,
+                lambda key, r: create_delta_chunk({key[1]: r}, model=model_name, index=key[0]),
+            ):
+                yield sse
+            for sse in _drain(
+                tool_bufs,
+                finished,
+                lambda key, r: create_delta_chunk(
+                    {"tool_calls": [{"index": key[1], "function": {"arguments": r}}]},
+                    model=model_name,
+                    index=key[0],
+                ),
+            ):
+                yield sse
+            for sse in _drain(
+                func_bufs,
+                finished,
+                lambda idx, r: create_delta_chunk(
+                    {"function_call": {"arguments": r}}, model=model_name, index=idx
+                ),
+            ):
+                yield sse
 
         except Exception:
             logger.exception("Error during streaming")
@@ -289,20 +301,18 @@ async def stream_text_response(
 
                 yield format_sse_event(json.dumps(chunk_dict))
 
-            # Stream ended without a finish chunk for some choice: emit whatever
-            # the buffers still hold so no restored text is silently dropped.
-            for idx, buf in buffers.items():
-                if idx in flushed:
-                    continue
-                remaining = buf.flush()
-                if remaining:
-                    yield format_sse_event(json.dumps({
-                        "object": "text_completion",
-                        "model": model_name,
-                        "choices": [
-                            {"index": idx, "text": remaining, "finish_reason": None}
-                        ],
-                    }))
+            # Stream ended without a finish chunk for some choice: emit whatever the
+            # buffers still hold so no restored text is silently dropped.
+            for sse in _drain(
+                buffers,
+                flushed,
+                lambda idx, r: {
+                    "object": "text_completion",
+                    "model": model_name,
+                    "choices": [{"index": idx, "text": r, "finish_reason": None}],
+                },
+            ):
+                yield sse
 
         except Exception:
             logger.exception("Error during text-completion streaming")

From 9524757baa030f7be4a99160dcf5fb2be4f8e92b Mon Sep 17 00:00:00 2001
From: crp4222 <121295348+crp4222@users.noreply.github.com>
Date: Thu, 9 Jul 2026 16:46:33 +0100
Subject: [PATCH 3/4] docs: clarify OpenWebUI connection and the two keys

---
 README.md | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 6ac332e..9d23189 100644
--- a/README.md
+++ b/README.md
@@ -224,11 +224,11 @@ python -m privaite --reload
 
 ### 4. Connect
 
-Point any OpenAI-compatible client to `http://localhost:8400/v1` with your proxy API key. Ready-to-run client snippets (curl, Python, Node) are in [`examples/`](examples/).
+Point any OpenAI-compatible client at `http://localhost:8400/v1` and authenticate with your `PRIVAITE_API_KEYS` value (not your provider key). Ready-to-run client snippets (curl, Python, Node) are in [`examples/`](examples/).
 
-**OpenWebUI (Docker):** Admin → Settings → Connections → OpenAI API:
-- URL: `http://host.docker.internal:8400/v1`
-- Key: your `PRIVAITE_API_KEYS` value
+**Open WebUI → Admin → Settings → Connections → OpenAI API:**
+- URL: `http://localhost:8400/v1`, or `http://host.docker.internal:8400/v1` if Open WebUI itself runs in Docker
+- Key: your `PRIVAITE_API_KEYS` value. Your real OpenAI key never goes here; it stays in the PrivAiTe container (see [Docker](#docker)).
 
 If you would rather not run a separate proxy, there is also an in-process Open
 WebUI filter (see [Open WebUI filter](#open-webui-filter) below).
@@ -247,6 +247,11 @@ docker run -d -p 8400:8400 \
   ghcr.io/crp4222/privaite
 ```
 
+Two keys, two roles, like any proxy: `PRIVAITE_API_KEYS` is what your client (Open
+WebUI, your app) sends to PrivAiTe; `OPENAI_API_KEY` is your real OpenAI key that
+PrivAiTe uses upstream. The provider key stays in the container and never reaches
+your client, which is the whole point.
+
 That exposes `gpt-4o-mini` and `gpt-4o`. For any other provider (Ollama, Azure, a
 self-hosted endpoint, or your own LiteLLM proxy), mount a config instead:
 

From a73f3cb713caa173b37ae2b359c16ebed2f410d8 Mon Sep 17 00:00:00 2001
From: crp4222 <121295348+crp4222@users.noreply.github.com>
Date: Thu, 9 Jul 2026 16:46:34 +0100
Subject: [PATCH 4/4] release: 0.3.0

---
 CHANGELOG.md         | 22 +++++++++++-----------
 privaite/__init__.py |  2 +-
 pyproject.toml       |  2 +-
 3 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4be0bde..39b3de9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -6,17 +6,17 @@ All notable changes to this project are documented here. The format follows
 
 ## [Unreleased]
 
-### Fixed
-- **Open WebUI filter: restore PII on Open WebUI >= 0.10** (filter bumped to
-  0.1.6). Since 0.10 the assistant reply is stored as structured `output` items
-  (`{type, content: [{type: "output_text", text}]}`) with `message["content"]`
-  empty, so the filter's `outlet` restore was a silent no-op and users saw
-  placeholders (`<PERSON_1>`) in the reply. `outlet` now also restores the text
-  of `output` items (message and reasoning), returning a new object only when a
-  value changed so Open WebUI persists it. The older `content` path still works
-  for Open WebUI < 0.10. Standalone proxy, CLI and the LiteLLM guardrail were
-  never affected (they never see the `output`-items shape). Hub listing needs a
-  manual republish to ship 0.1.6.
+## [0.3.0] - 2026-07-09
+
+### Added
+- Official Docker image at `ghcr.io/crp4222/privaite` (multi-arch, detection model
+  baked in, runs offline). Drop-in for OpenAI:
+  `docker run -e PRIVAITE_API_KEYS=change-me -e OPENAI_API_KEY=sk-... ghcr.io/crp4222/privaite`.
+
+### Changed
+- Internal refactoring for a cleaner, less-duplicated codebase. No change in
+  behavior or detection.
+- Clearer Open WebUI / Docker connection docs (which URL, which key).
 
 ## [0.2.13] - 2026-07-03
 
diff --git a/privaite/__init__.py b/privaite/__init__.py
index 11ef092..493f741 100644
--- a/privaite/__init__.py
+++ b/privaite/__init__.py
@@ -1 +1 @@
-__version__ = "0.2.13"
+__version__ = "0.3.0"
diff --git a/pyproject.toml b/pyproject.toml
index 7befae5..fa4b1fc 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "privaite"
-version = "0.2.13"
+version = "0.3.0"
 description = "Drop-in self-hosted LLM proxy that reversibly redacts PII before LLM calls, including tool-call arguments and multimodal content"
 readme = "README.md"
 requires-python = ">=3.11"