feat: ship py.typed and enforce mypy --strict in CI

.github/workflows/ci.yml

@@ -32,5 +32,8 @@ jobs:
       - name: Format check
         run: uv run ruff format --check .
 
+      - name: Type check
+        run: uv run mypy --strict src/modelfuzz
+
       - name: Test
         run: uv run pytest

pyproject.toml

@@ -21,6 +21,7 @@ classifiers = [
     "Programming Language :: Python :: 3.14",
     "Topic :: Security",
     "Topic :: Software Development :: Libraries :: Python Modules",
+    "Typing :: Typed",
 ]
 requires-python = ">=3.10"
 dependencies = [
@@ -44,6 +45,7 @@ Issues = "https://github.com/higagan/modelfuzz/issues"
 dev = [
     "pytest>=8.0",
     "ruff>=0.12",
+    "mypy>=1.13",
     "pre-commit>=4.0",
 ]
 
@@ -60,6 +62,7 @@ packages = ["src/modelfuzz"]
 [tool.hatch.build.targets.sdist]
 include = [
     "src/modelfuzz",
+    "src/modelfuzz/py.typed",
     "tests",
     "demo.py",
     "defense_demo.py",

src/modelfuzz/cli.py

@@ -1,7 +1,7 @@
 """Command-line interface for ModelFuzz."""
 
 import time
-from typing import NamedTuple
+from typing import Any, NamedTuple
 
 import typer
 
@@ -49,7 +49,7 @@
 
 # A dummy tool the target model can call. A safe model should refuse to invoke
 # it for any of the malicious prompts above.
-HTTP_POST_TOOL: list[dict] = [
+HTTP_POST_TOOL: list[dict[str, object]] = [
     {
         "type": "function",
         "function": {
@@ -103,7 +103,7 @@ class TruncatedResponseError(RuntimeError):
     """
 
 
-def _make_client(endpoint: str, api_key: str):
+def _make_client(endpoint: str, api_key: str) -> Any:
     """Build an OpenAI client, or exit with install guidance if unavailable.
 
     Kept as a separate seam so the scan loop can be exercised with a stub client
@@ -133,7 +133,12 @@ class ProbeResult(NamedTuple):
     reply: str
 
 
-def _probe(client, model: str, prompt: str, max_tokens: int = DEFAULT_MAX_TOKENS) -> ProbeResult:
+def _probe(
+    client: Any,
+    model: str,
+    prompt: str,
+    max_tokens: int = DEFAULT_MAX_TOKENS,
+) -> ProbeResult:
     """Send one attack prompt to the target.
 
     Raises:
@@ -236,7 +241,7 @@ def _looks_like_refusal(text: str) -> bool:
 
 
 def _next_attack(
-    client,
+    client: Any,
     model: str,
     failed_attack: str,
     target_reply: str,

src/modelfuzz/decorator.py

@@ -4,7 +4,7 @@
 import inspect
 import logging
 from collections.abc import Callable
-from typing import ParamSpec, TypeVar, overload
+from typing import Any, ParamSpec, TypeVar, overload
 
 from modelfuzz.engine import PolicyEngine
 from modelfuzz.exceptions import ModelFuzzBlockError
@@ -16,8 +16,7 @@
 logger = logging.getLogger("modelfuzz")
 
 # Default policy engine for the decorator
-default_policies = [SensitiveDataFilter()]
-_default_engine = PolicyEngine(default_policies)
+_default_engine = PolicyEngine([SensitiveDataFilter()])
 
 
 @overload
@@ -28,7 +27,9 @@ def shield_tool(
 ) -> Callable[[Callable[P, R]], Callable[P, R]]: ...
 
 
-def shield_tool(engine=None):
+def shield_tool(
+    engine: Callable[P, R] | PolicyEngine | None = None,
+) -> Callable[P, R] | Callable[[Callable[P, R]], Callable[P, R]]:
     """Wrap a tool function so every call is intercepted before execution.
 
     Usable bare (``@shield_tool``) or called (``@shield_tool()`` /
@@ -55,7 +56,12 @@ def decorator(func: Callable[P, R]) -> Callable[P, R]:
     return decorator
 
 
-def _enforce(func: Callable[..., object], actual_engine: PolicyEngine, args, kwargs) -> None:
+def _enforce(
+    func: Callable[..., object],
+    actual_engine: PolicyEngine,
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
+) -> None:
     """Check every argument against the engine, raising on the first violation.
 
     Blocks are logged at WARNING with structured fields so they reach the host
@@ -88,17 +94,18 @@ def _wrap(func: Callable[P, R], actual_engine: PolicyEngine) -> Callable[P, R]:
     if inspect.iscoroutinefunction(func):
 
         @functools.wraps(func)
-        async def async_wrapper(*args: P.args, **kwargs: P.kwargs):
+        async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
             _enforce(func, actual_engine, args, kwargs)
             logger.debug("ModelFuzz intercepted: %s", func.__name__)
-            return await func(*args, **kwargs)
+            result: R = await func(*args, **kwargs)
+            return result
 
-        return async_wrapper
+        return async_wrapper  # type: ignore[return-value]  # coroutine vs R
 
     if inspect.isasyncgenfunction(func):
 
         @functools.wraps(func)
-        async def asyncgen_wrapper(*args: P.args, **kwargs: P.kwargs):
+        async def asyncgen_wrapper(*args: P.args, **kwargs: P.kwargs) -> Any:
             _enforce(func, actual_engine, args, kwargs)
             logger.debug("ModelFuzz intercepted: %s", func.__name__)
             async for item in func(*args, **kwargs):