stackone-defender 0.7.2__tar.gz → 0.7.4__tar.gz

stackone_defender-0.7.4/.release-please-manifest.json

@@ -0,0 +1 @@
+{".":"0.7.4"}

{stackone_defender-0.7.2 → stackone_defender-0.7.4}/CHANGELOG.md

@@ -1,5 +1,31 @@
 # Changelog
 
+## [0.7.4](https://github.com/StackOneHQ/defender-py/compare/stackone-defender-v0.7.3...stackone-defender-v0.7.4) (2026-08-13)
+
+
+### Features
+
+* **ENG-1761:** port Tier 2 buckets, dedupe, and cost telemetry from TS 0.7.4 ([b92b9f6](https://github.com/StackOneHQ/defender-py/commit/b92b9f661ef927b27542a664d92477e039279758))
+* **ENG-1761:** port Tier 2 fixed-width buckets, dedupe, cost telemetry from TS 0.7.4 ([1e63d09](https://github.com/StackOneHQ/defender-py/commit/1e63d0999d003348c17d6700bf6a03b65a7b2242))
+
+
+### Bug Fixes
+
+* **ENG-1761:** address review — cold_load on failure paths + stronger tests ([e897786](https://github.com/StackOneHQ/defender-py/commit/e897786ff013b95b7b74eecd8df554ee978dd98b))
+
+## [0.7.3](https://github.com/StackOneHQ/defender-py/compare/stackone-defender-v0.7.2...stackone-defender-v0.7.3) (2026-07-22)
+
+
+### Features
+
+* **ENG-1279:** port Tier 3 score-threshold (decide on score, not the model's word) ([77c8a6a](https://github.com/StackOneHQ/defender-py/commit/77c8a6a38336ddd797f5691cedf4670f15f575fb))
+
+
+### Bug Fixes
+
+* **ENG-1279:** Tier 3 score threshold + ENG-1296 Tier 2/base64/SFE fixes ([3f2b9cd](https://github.com/StackOneHQ/defender-py/commit/3f2b9cd1f7dd82f19e468076bcb3196ad22bf83e))
+* **ENG-1296:** port TS parity fixes for Tier 2 truncation, base64, SFE recursion ([4fa8cda](https://github.com/StackOneHQ/defender-py/commit/4fa8cda445d8c08e30e74d6053e2eaf26026fed2))
+
 ## [0.7.2](https://github.com/StackOneHQ/stackone-defender/compare/stackone-defender-v0.7.1...stackone-defender-v0.7.2) (2026-06-30)
 
 

stackone_defender-0.7.4/LICENSE

@@ -0,0 +1,190 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but not
+      limited to compiled object code, generated documentation, and
+      conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to the Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by the Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding any notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   Copyright 2024 StackOne
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

{stackone_defender-0.7.2 → stackone_defender-0.7.4}/PKG-INFO

@@ -1,14 +1,14 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: stackone-defender
-Version: 0.7.2
+Version: 0.7.4
 Summary: Indirect prompt injection defense for AI agents using tool calls
 Project-URL: Homepage, https://github.com/StackOneHQ/stackone-defender
 Project-URL: Repository, https://github.com/StackOneHQ/stackone-defender
 Author-email: StackOne <support@stackone.com>
-License: Apache-2.0
+License-Expression: Apache-2.0
+License-File: LICENSE
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Developers
-Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
@@ -146,6 +146,47 @@ Packed-chunk MiniLM classifier (int8 ONNX ~22 MB, bundled):
 | jayavibhav (adversarial) | 0.9717 | ~1k |
 | **Average** | **0.9079** | ~25k |
 
+### Optional Tier 3 — LLM adjudication (consumer-supplied)
+
+Authoritative LLM-based classification for the cases Tier 2 finds ambiguous. The package ships **only** the orchestration and the `Tier3Provider` interface — the actual model endpoint (a hosted LLM, OpenAI, an internal inference service) lives in your code, keeping proprietary models and credentials out of the package.
+
+Two modes, selected via `defender_mode`:
+- **`"cascade"`** (default): Tier 1 → Tier 2 → Tier 3, with Tier 3 invoked only when the Tier 2 effective score falls in the gray band (default `[0.3, 0.85)`). The Tier 3 verdict overrides Tier 2 on the escalated chunk — a `block` forces a block, an `allow` rescues it. Outside the band defender skips the round trip.
+- **`"tier3_only"`**: skip Tier 2; the block/allow decision is the Tier 3 verdict alone. Tier 1 still sanitizes the returned `sanitized` payload.
+
+Register a provider once at startup, then opt in per instance:
+
+```python
+from typing import Any
+
+from stackone_defender import create_prompt_defense, set_default_tier3_provider
+from stackone_defender.types import Tier3Verdict
+
+class MyProvider:
+    def classify(self, text: str, *, ctx: dict[str, Any] | None = None) -> Tier3Verdict:
+        # Call your LLM endpoint here (sync or awaitable). ctx["toolName"] is available.
+        verdict = call_my_llm(text, tool_name=(ctx or {}).get("toolName"))
+        return Tier3Verdict(decision="block" if verdict.block else "allow", score=verdict.p_block)
+
+set_default_tier3_provider(MyProvider())
+
+defense = create_prompt_defense(
+    block_high_risk=True,
+    enable_tier3=True,
+    defender_mode="cascade",  # or "tier3_only"
+    tier3={
+        "provider": MyProvider(),                          # optional: overrides the registry for this instance
+        "escalation_band": {"lower": 0.3, "upper": 0.85},  # cascade gray band; [lower, upper), defaults shown
+        "max_text_length": 10000,                          # caps text passed to the provider
+        "block_threshold": 0.622,                          # optional; decide on score, not the model's word
+    },
+)
+```
+
+**Choosing the operating point (`block_threshold`).** By default the model's `decision` word is authoritative — but that word is its argmax, an implicit 0.5 cut that moves on its own whenever the model is retrained. Set `tier3.block_threshold` to decide on `verdict.score` (P(block)) instead: the cut becomes an explicit config value — raise it to trade recall for fewer false positives, lower it for the reverse; `0.5` reproduces argmax. It requires a provider that reports `score` as P(block) (not "confidence in whichever decision I made" — those invert on allows). If `score` is missing or outside `[0, 1]` the verdict's `decision` is used instead and defender warns once, so a provider that cannot report a score degrades to the default rather than failing.
+
+**Fail-open semantics.** A provider error or timeout records a `skip_reason` on `result.tier3`; in cascade defender falls back to the Tier 2 decision, in `tier3_only` it allows the request. `enable_tier3=True` with no registered provider falls back to the standard Tier 1 + Tier 2 cascade and logs one warning per instance — Tier 3 misconfiguration never silently disables defense. `result.tier3` carries the verdict (a `Tier3Verdict`, or a `Tier3Skip` when the provider ran but returned nothing usable) when Tier 3 runs, and is `None` when it doesn't.
+
 ### `allowed` vs `risk_level`
 
 - Use **`allowed`** for gating when `block_high_risk=True`: `False` means do not pass `sanitized` to the model as-is.

{stackone_defender-0.7.2 → stackone_defender-0.7.4}/README.md

@@ -120,6 +120,47 @@ Packed-chunk MiniLM classifier (int8 ONNX ~22 MB, bundled):
 | jayavibhav (adversarial) | 0.9717 | ~1k |
 | **Average** | **0.9079** | ~25k |
 
+### Optional Tier 3 — LLM adjudication (consumer-supplied)
+
+Authoritative LLM-based classification for the cases Tier 2 finds ambiguous. The package ships **only** the orchestration and the `Tier3Provider` interface — the actual model endpoint (a hosted LLM, OpenAI, an internal inference service) lives in your code, keeping proprietary models and credentials out of the package.
+
+Two modes, selected via `defender_mode`:
+- **`"cascade"`** (default): Tier 1 → Tier 2 → Tier 3, with Tier 3 invoked only when the Tier 2 effective score falls in the gray band (default `[0.3, 0.85)`). The Tier 3 verdict overrides Tier 2 on the escalated chunk — a `block` forces a block, an `allow` rescues it. Outside the band defender skips the round trip.
+- **`"tier3_only"`**: skip Tier 2; the block/allow decision is the Tier 3 verdict alone. Tier 1 still sanitizes the returned `sanitized` payload.
+
+Register a provider once at startup, then opt in per instance:
+
+```python
+from typing import Any
+
+from stackone_defender import create_prompt_defense, set_default_tier3_provider
+from stackone_defender.types import Tier3Verdict
+
+class MyProvider:
+    def classify(self, text: str, *, ctx: dict[str, Any] | None = None) -> Tier3Verdict:
+        # Call your LLM endpoint here (sync or awaitable). ctx["toolName"] is available.
+        verdict = call_my_llm(text, tool_name=(ctx or {}).get("toolName"))
+        return Tier3Verdict(decision="block" if verdict.block else "allow", score=verdict.p_block)
+
+set_default_tier3_provider(MyProvider())
+
+defense = create_prompt_defense(
+    block_high_risk=True,
+    enable_tier3=True,
+    defender_mode="cascade",  # or "tier3_only"
+    tier3={
+        "provider": MyProvider(),                          # optional: overrides the registry for this instance
+        "escalation_band": {"lower": 0.3, "upper": 0.85},  # cascade gray band; [lower, upper), defaults shown
+        "max_text_length": 10000,                          # caps text passed to the provider
+        "block_threshold": 0.622,                          # optional; decide on score, not the model's word
+    },
+)
+```
+
+**Choosing the operating point (`block_threshold`).** By default the model's `decision` word is authoritative — but that word is its argmax, an implicit 0.5 cut that moves on its own whenever the model is retrained. Set `tier3.block_threshold` to decide on `verdict.score` (P(block)) instead: the cut becomes an explicit config value — raise it to trade recall for fewer false positives, lower it for the reverse; `0.5` reproduces argmax. It requires a provider that reports `score` as P(block) (not "confidence in whichever decision I made" — those invert on allows). If `score` is missing or outside `[0, 1]` the verdict's `decision` is used instead and defender warns once, so a provider that cannot report a score degrades to the default rather than failing.
+
+**Fail-open semantics.** A provider error or timeout records a `skip_reason` on `result.tier3`; in cascade defender falls back to the Tier 2 decision, in `tier3_only` it allows the request. `enable_tier3=True` with no registered provider falls back to the standard Tier 1 + Tier 2 cascade and logs one warning per instance — Tier 3 misconfiguration never silently disables defense. `result.tier3` carries the verdict (a `Tier3Verdict`, or a `Tier3Skip` when the provider ran but returned nothing usable) when Tier 3 runs, and is `None` when it doesn't.
+
 ### `allowed` vs `risk_level`
 
 - Use **`allowed`** for gating when `block_high_risk=True`: `False` means do not pass `sanitized` to the model as-is.

{stackone_defender-0.7.2 → stackone_defender-0.7.4}/pyproject.toml

@@ -1,15 +1,15 @@
 [project]
 name = "stackone-defender"
-version = "0.7.2"
+version = "0.7.4"
 description = "Indirect prompt injection defense for AI agents using tool calls"
 readme = "README.md"
 requires-python = ">=3.11"
-license = { text = "Apache-2.0" }
+license = "Apache-2.0"
+license-files = ["LICENSE"]
 authors = [{ name = "StackOne", email = "support@stackone.com" }]
 classifiers = [
   "Development Status :: 4 - Beta",
   "Intended Audience :: Developers",
-  "License :: OSI Approved :: Apache Software License",
   "Programming Language :: Python :: 3",
   "Programming Language :: Python :: 3.11",
   "Programming Language :: Python :: 3.12",
@@ -32,6 +32,8 @@ sfe = ["fasttext-ng>=0.9.3"]
 [dependency-groups]
 dev = [
   "pytest>=8.0",
+  "pytest-cov>=5.0",
+  "ruff>=0.9",
   "onnxruntime>=1.16.0",
   "tokenizers>=0.15.0",
   "numpy>=1.24.0",
@@ -45,5 +47,26 @@ build-backend = "hatchling.build"
 [tool.hatch.build.targets.wheel]
 packages = ["src/stackone_defender"]
 
+[tool.hatch.build.targets.sdist]
+exclude = ["/models", "/.github"]  # legacy model + CI, dev-only
+
+[tool.ruff]
+target-version = "py311"
+line-length = 120
+src = ["src", "tests"]
+
+[tool.ruff.lint]
+select = ["E", "W", "F", "I", "N", "UP", "B", "C4", "SIM", "RUF"]
+# E501/RUF003 by design; SIM*/UP007 deferred to cleanup.
+ignore = ["E501", "RUF003", "SIM102", "SIM103", "SIM108", "SIM110", "SIM118", "UP007"]
+
+[tool.ruff.lint.per-file-ignores]
+"tests/*" = ["RUF012", "N814", "F841", "F811"]
+
 [tool.pytest.ini_options]
 testpaths = ["tests"]
+addopts = "--strict-markers --strict-config -ra"
+markers = [
+  "model: requires the bundled ONNX model + onnxruntime extra",
+  "sfe: requires the fasttext extra",
+]

{stackone_defender-0.7.2 → stackone_defender-0.7.4}/src/stackone_defender/classifiers/onnx_classifier.py

@@ -11,13 +11,24 @@ from __future__ import annotations
 import logging
 import math
 import threading
+from dataclasses import dataclass
 from pathlib import Path
-from typing import Literal
+from typing import Literal, cast
 
 _logger = logging.getLogger(__name__)
 
+
+@dataclass
+class BatchTokenStats:
+    """Padding accounting accumulated across a batched classification (#6 telemetry)."""
+
+    real_tokens: int = 0  # sum of real (non-pad) tokens across all chunks
+    padded_tokens: int = 0  # tokens actually run through ONNX, including padding
+
 # Shared across all OnnxClassifier instances (keyed by resolved model dir path).
-_session_cache: dict[str, tuple[object, object]] = {}
+# Tuple is (session, tokenizer, count_tokenizer); the count tokenizer is
+# non-truncating so count_tokens reports true length (see _load_model).
+_session_cache: dict[str, tuple[object, object, object]] = {}
 _registry_lock = threading.Lock()
 _load_locks: dict[str, threading.Lock] = {}
 
@@ -59,11 +70,17 @@ class OnnxClassifier:
     """
 
     _MAX_BATCH_CHUNK = 32
+    # Fixed pad widths (ascending). A string pads to the smallest bucket >= its
+    # token length, so its padded length — and thus its quantized score — is a
+    # function of the string alone, not its batch neighbours (deterministic).
+    _PAD_BUCKETS = (32, 64, 128, 256)
 
     def __init__(self, model_path: str | None = None, temperature_t: float | None = None):
         self._model_path = model_path or get_default_model_path()
         self._session = None
         self._tokenizer = None
+        # Non-truncating tokenizer used only by count_tokens (see _load_model).
+        self._count_tokenizer = None
         self._max_length = 256
         self._load_failed = False
         # Output mode is detected lazily from the logits shape on the first
@@ -114,13 +131,13 @@ class OnnxClassifier:
         cache_key = str(Path(self._model_path).resolve())
         cached = _session_cache.get(cache_key)
         if cached:
-            self._session, self._tokenizer = cached
+            self._session, self._tokenizer, self._count_tokenizer = cached
             return
 
         with _lock_for_cache_key(cache_key):
             cached = _session_cache.get(cache_key)
             if cached:
-                self._session, self._tokenizer = cached
+                self._session, self._tokenizer, self._count_tokenizer = cached
                 return
 
             try:
@@ -138,7 +155,18 @@ class OnnxClassifier:
                 tokenizer_path = str(Path(self._model_path) / "tokenizer.json")
                 self._tokenizer = Tokenizer.from_file(tokenizer_path)
                 self._tokenizer.enable_truncation(max_length=self._max_length)
-                self._tokenizer.enable_padding(length=self._max_length)
+                # No fixed padding: classify_batch_pair pads each batch per bucket
+                # width (see _PAD_BUCKETS), so a string's padded length depends only
+                # on itself. Single classify runs at the string's real length.
+                self._tokenizer.no_padding()
+
+                # tokenizer.json bakes in truncation+padding, so disable both
+                # here to count true length rather than a value capped at
+                # max_length (which would make every long payload look like a
+                # single chunk and drop everything past the cap).
+                self._count_tokenizer = Tokenizer.from_file(tokenizer_path)
+                self._count_tokenizer.no_truncation()
+                self._count_tokenizer.no_padding()
 
                 onnx_path = str(Path(self._model_path) / "model_quantized.onnx")
                 self._session = ort.InferenceSession(onnx_path)
@@ -146,7 +174,7 @@ class OnnxClassifier:
                 _logger.warning("[defender] ONNX model failed to load: %s", e)
                 raise
 
-            _session_cache[cache_key] = (self._session, self._tokenizer)
+            _session_cache[cache_key] = (self._session, self._tokenizer, self._count_tokenizer)
 
     # ------------------------------------------------------------------
     # Inference
@@ -187,36 +215,77 @@ class OnnxClassifier:
         main_logit = float(row[0]) if hasattr(row, "__len__") and len(row) > 0 else float(row)
         return _sigmoid(main_logit / t), None
 
-    def classify_batch(self, texts: list[str]) -> list[float]:
+    def classify_batch(self, texts: list[str], stats: BatchTokenStats | None = None) -> list[float]:
         """Classify multiple texts; returns main-head scores only.
 
         Back-compat wrapper around :meth:`classify_batch_pair`.
         """
-        return [main for main, _ in self.classify_batch_pair(texts)]
+        return [main for main, _ in self.classify_batch_pair(texts, stats)]
 
-    def classify_batch_pair(self, texts: list[str]) -> list[tuple[float, float | None]]:
+    def classify_batch_pair(
+        self, texts: list[str], stats: BatchTokenStats | None = None
+    ) -> list[tuple[float, float | None]]:
         """Classify multiple texts, returning ``(main, aux)`` per row.
 
-        Aux is ``None`` per-row for single-head models. Chunks the input to
-        bound native memory; the attention matrix is ``O(chunk * seq_len^2)``,
-        and for MiniLM (``max_length=256``) a chunk of 32 keeps memory
-        under ~50MB per call.
+        Aux is ``None`` per-row for single-head models. Groups strings into fixed
+        pad-width buckets by token length, batches within each bucket, and pads
+        every row to the BUCKET width (not the batch max). This bounds padding
+        waste and makes a string's padded length — hence its quantized score —
+        depend only on itself, not its batch neighbours. Chunks of 32 bound
+        native memory (attention is ``O(chunk * seq_len^2)``). Results scatter
+        back to input order. ``stats`` accumulates padding counts when provided.
         """
         if not texts:
             return []
         self._ensure_loaded()
-        all_pairs: list[tuple[float, float | None]] = []
-        for offset in range(0, len(texts), self._MAX_BATCH_CHUNK):
-            chunk = texts[offset : offset + self._MAX_BATCH_CHUNK]
-            all_pairs.extend(self._classify_batch_chunk_pair(chunk))
-        return all_pairs
 
-    def _classify_batch_chunk_pair(self, texts: list[str]) -> list[tuple[float, float | None]]:
+        encodings = self._tokenizer.encode_batch(texts)
+        last_bucket = self._PAD_BUCKETS[-1]
+        buckets: dict[int, list[int]] = {}
+        for i, enc in enumerate(encodings):
+            length = len(enc.ids)
+            width = next((b for b in self._PAD_BUCKETS if length <= b), last_bucket)
+            buckets.setdefault(width, []).append(i)
+
+        pairs: list[tuple[float, float | None] | None] = [None] * len(texts)
+        for width in self._PAD_BUCKETS:
+            bucket_idxs = buckets.get(width)
+            if not bucket_idxs:
+                continue
+            for offset in range(0, len(bucket_idxs), self._MAX_BATCH_CHUNK):
+                idxs = bucket_idxs[offset : offset + self._MAX_BATCH_CHUNK]
+                chunk = [encodings[i] for i in idxs]
+                if stats is not None:
+                    # Every row pads to the bucket width; real_tokens is useful work.
+                    stats.padded_tokens += len(chunk) * width
+                    for enc in chunk:
+                        stats.real_tokens += len(enc.ids)
+                chunk_pairs = self._classify_batch_chunk_pair(chunk, width)
+                for k, orig_idx in enumerate(idxs):
+                    pairs[orig_idx] = chunk_pairs[k]
+
+        return cast(list[tuple[float, float | None]], pairs)
+
+    def _classify_batch_chunk_pair(
+        self, encodings: list, pad_to: int | None = None
+    ) -> list[tuple[float, float | None]]:
+        """Run one pre-tokenized chunk through ONNX, padding rows to ``pad_to``.
+
+        Rows are padded to at least the chunk's real max, so a caller can never
+        under-pad. Inputs are pre-tokenized so ``classify_batch_pair`` can
+        length-bucket without re-tokenizing.
+        """
         import numpy as np
 
-        encodings = self._tokenizer.encode_batch(texts)
-        input_ids = np.array([e.ids for e in encodings], dtype=np.int64)
-        attention_mask = np.array([e.attention_mask for e in encodings], dtype=np.int64)
+        actual_max = max(len(e.ids) for e in encodings)
+        max_len = max(pad_to, actual_max) if pad_to is not None else actual_max
+        batch_size = len(encodings)
+        input_ids = np.zeros((batch_size, max_len), dtype=np.int64)
+        attention_mask = np.zeros((batch_size, max_len), dtype=np.int64)
+        for i, enc in enumerate(encodings):
+            n = len(enc.ids)
+            input_ids[i, :n] = enc.ids
+            attention_mask[i, :n] = enc.attention_mask
 
         results = self._session.run(None, {"input_ids": input_ids, "attention_mask": attention_mask})
         logits = results[0]
@@ -225,12 +294,12 @@ class OnnxClassifier:
         t = self._temperature_t
         pairs: list[tuple[float, float | None]] = []
         if self._output_mode == "multi":
-            for i in range(len(texts)):
+            for i in range(batch_size):
                 main = _sigmoid(float(logits[i][0]) / t)
                 aux = _sigmoid(float(logits[i][1]) / t)
                 pairs.append((main, aux))
         else:
-            for i in range(len(texts)):
+            for i in range(batch_size):
                 row = logits[i]
                 # ``row`` may be a scalar (shape ``[batch]``) or 1-vector.
                 main_logit = float(row[0]) if hasattr(row, "__len__") and len(row) > 0 else float(row)
@@ -258,9 +327,9 @@ class OnnxClassifier:
 
     def count_tokens(self, text: str) -> int:
         self._ensure_loaded()
-        encoding = self._tokenizer.encode(text)
-        # Padding is enabled at a fixed length; count only real (attended) tokens.
-        return int(sum(encoding.attention_mask))
+        # Non-truncating count, including special tokens ([CLS]/[SEP]) to match
+        # the TS countTokens; the inference tokenizer would cap this at max_length.
+        return len(self._count_tokenizer.encode(text).ids)
 
     def get_max_length(self) -> int:
         return self._max_length

{stackone_defender-0.7.2 → stackone_defender-0.7.4}/src/stackone_defender/classifiers/tier2_classifier.py

@@ -17,7 +17,7 @@ from typing import Any
 
 from ..types import MultiheadConfig, RiskLevel, Tier2Result
 from ..utils.boundary import strip_boundary_patterns
-from .onnx_classifier import OnnxClassifier, get_default_model_path
+from .onnx_classifier import BatchTokenStats, OnnxClassifier, get_default_model_path
 
 _logger = logging.getLogger(__name__)
 
@@ -343,14 +343,16 @@ class Tier2Classifier:
             return {"chunks": [], "skipped": True, "skip_reason": "No classifiable sentences"}
         return {"chunks": self._pack_sentences(sentences, max_content_tokens), "skipped": False}
 
-    def classify_chunks_batch(self, chunks: list[str]) -> list[float]:
+    def classify_chunks_batch(self, chunks: list[str], stats: BatchTokenStats | None = None) -> list[float]:
         """Single-head batch classify. Returns main-head scores only."""
         if not chunks:
             return []
         self._onnx.warmup()
-        return self._onnx.classify_batch(chunks)
+        return self._onnx.classify_batch(chunks, stats)
 
-    def classify_chunks_batch_pair(self, chunks: list[str]) -> list[tuple[float, float | None]]:
+    def classify_chunks_batch_pair(
+        self, chunks: list[str], stats: BatchTokenStats | None = None
+    ) -> list[tuple[float, float | None]]:
         """Multi-head variant. Returns ``(main, aux)`` per chunk.
 
         Aux is ``None`` per-row for single-head models. Callers in the
@@ -359,7 +361,7 @@ class Tier2Classifier:
         if not chunks:
             return []
         self._onnx.warmup()
-        return self._onnx.classify_batch_pair(chunks)
+        return self._onnx.classify_batch_pair(chunks, stats)
 
     # ------------------------------------------------------------------
     # Helpers / introspection