diff --git a/.claude/settings.json b/.claude/settings.json index 9b09ca7..0b077ee 100644 --- a/.claude/settings.json +++ b/.claude/settings.json @@ -17,6 +17,8 @@ "Bash(docker compose -f docker/compose.yaml run --rm --no-deps php-example)", "Bash(docker compose -f docker/compose.yaml run --rm --no-deps mcp-fetcher)", "Bash(docker compose -f docker/compose.yaml run --rm --no-deps mcp-tests)", + "Bash(docker compose -f docker/compose.yaml run --rm --no-deps py-dev-fetcher)", + "Bash(docker compose -f docker/compose.yaml run --rm --no-deps py-analysis)", "Bash(docker compose -f docker/compose.yaml run --rm --no-deps composer-install)", "Bash(docker compose -f docker/compose.yaml run --rm --no-deps php-fetcher)", "Bash(docker compose -f docker/compose.yaml run --rm --no-deps php-analysis)", @@ -75,6 +77,8 @@ "PowerShell(docker compose -f docker/compose.yaml run --rm --no-deps php-example)", "PowerShell(docker compose -f docker/compose.yaml run --rm --no-deps mcp-fetcher)", "PowerShell(docker compose -f docker/compose.yaml run --rm --no-deps mcp-tests)", + "PowerShell(docker compose -f docker/compose.yaml run --rm --no-deps py-dev-fetcher)", + "PowerShell(docker compose -f docker/compose.yaml run --rm --no-deps py-analysis)", "PowerShell(docker compose -f docker/compose.yaml run --rm --no-deps composer-install)", "PowerShell(docker compose -f docker/compose.yaml run --rm --no-deps php-fetcher)", "PowerShell(docker compose -f docker/compose.yaml run --rm --no-deps php-analysis)", diff --git a/.claude/skills/check/SKILL.md b/.claude/skills/check/SKILL.md index d2c5589..2e222a0 100644 --- a/.claude/skills/check/SKILL.md +++ b/.claude/skills/check/SKILL.md @@ -20,9 +20,11 @@ docker compose -f docker/compose.yaml up -d worker frankenphp fault-worker hosti docker compose -f docker/compose.yaml run --rm --no-deps php-integration docker compose -f docker/compose.yaml run --rm --no-deps integration docker compose -f docker/compose.yaml --profile test down -docker compose -f docker/compose.yaml --profile mcp --profile mcp-fetch --profile mcp-model --profile php-dev --profile php-frameworks --profile tls config +docker compose -f docker/compose.yaml --profile mcp --profile mcp-fetch --profile mcp-model --profile php-dev --profile py-dev --profile php-frameworks --profile tls config docker compose -f docker/compose.yaml run --rm --no-deps mcp-fetcher docker compose -f docker/compose.yaml run --rm --no-deps mcp-tests +docker compose -f docker/compose.yaml run --rm --no-deps py-dev-fetcher +docker compose -f docker/compose.yaml run --rm --no-deps py-analysis docker compose -f docker/compose.yaml run --rm --no-deps composer-install docker compose -f docker/compose.yaml run --rm --no-deps php-fetcher docker compose -f docker/compose.yaml run --rm --no-deps php-frameworks-fetcher @@ -33,7 +35,7 @@ docker compose -f docker/compose.yaml run --rm --no-deps laravel-tests docker compose -f docker/compose.yaml run --rm --no-deps symfony-tests docker compose -f docker/compose.yaml up -d tls-proxy docker compose -f docker/compose.yaml run --rm --no-deps tls-check -docker compose -f docker/compose.yaml --profile mcp --profile mcp-fetch --profile php-dev --profile php-frameworks --profile tls down +docker compose -f docker/compose.yaml --profile mcp --profile mcp-fetch --profile php-dev --profile py-dev --profile php-frameworks --profile tls down ``` The `composer-install` step is the consumer-install check: `composer validate --strict` on the library, then a @@ -42,6 +44,11 @@ repository in a git worktree whose git directory is outside the mount) with Pack `vendor/`, and autoloads `PhpAiBridge\Client`; the archive is built from the committed `HEAD`, so an uncommitted change is not what it tests, and it needs no network and no cache. +`py-dev-fetcher` downloads the mypy and ruff wheels `requirements-dev.lock` pins into `.cache/wheels`, with network, +and `py-analysis` installs them and the MCP server's wheels from `mcp-fetcher` into a tmpfs and runs, offline, +`ruff check` over `worker`, `mcp_server`, `scripts` and `tests` and `mypy --strict` over `worker` and `mcp_server`, +both configured in `pyproject.toml`. Fix a finding with annotations or a small guard, never a behaviour change. + `php-fetcher` installs the PHPStan and the PSR-7/17/18 interfaces that `composer.lock` pins into `.cache/php-vendor`, with network, and `php-analysis` runs PHPStan offline at level max over `src`, `examples` and `tests/php` with `phpstan.neon`; `phpstan-baseline.neon` holds the findings left on purpose, each with its diff --git a/.gitattributes b/.gitattributes index 0a1c626..df9e0d4 100644 --- a/.gitattributes +++ b/.gitattributes @@ -18,6 +18,9 @@ /mcp_server export-ignore /phpstan-baseline.neon export-ignore /phpstan.neon export-ignore +/pyproject.toml export-ignore +/requirements-dev.lock export-ignore +/requirements-dev.txt export-ignore /requirements-mcp.lock export-ignore /requirements-mcp.txt export-ignore /requirements-model.lock export-ignore diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 5aef097..b960ecd 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -69,11 +69,16 @@ jobs: - name: Prepare disposable cache run: mkdir -p .cache && chmod 777 .cache - name: Inspect sandbox configuration - run: docker compose -f docker/compose.yaml --profile mcp --profile mcp-fetch --profile mcp-model --profile php-dev --profile php-frameworks --profile tls config + run: docker compose -f docker/compose.yaml --profile mcp --profile mcp-fetch --profile mcp-model --profile php-dev --profile py-dev --profile php-frameworks --profile tls config - name: Acquire hash-pinned MCP wheels run: docker compose -f docker/compose.yaml run --rm --no-deps mcp-fetcher - name: Offline MCP server tests run: docker compose -f docker/compose.yaml run --rm --no-deps mcp-tests + - name: Acquire the pinned mypy and ruff + run: docker compose -f docker/compose.yaml run --rm --no-deps py-dev-fetcher + # After mcp-fetcher: mypy reads the MCP SDK's types from its wheels. + - name: mypy --strict and ruff over the Python code, offline + run: docker compose -f docker/compose.yaml run --rm --no-deps py-analysis - name: Consumer Composer install of the library run: docker compose -f docker/compose.yaml run --rm --no-deps -e BRIDGE_REQUIRE_ARCHIVE=1 composer-install - name: Acquire the locked PHPStan @@ -97,7 +102,7 @@ jobs: docker compose -f docker/compose.yaml run --rm --no-deps tls-check - name: Stop test containers if: always() - run: docker compose -f docker/compose.yaml --profile mcp --profile mcp-fetch --profile php-dev --profile php-frameworks --profile tls down + run: docker compose -f docker/compose.yaml --profile mcp --profile mcp-fetch --profile php-dev --profile py-dev --profile php-frameworks --profile tls down # The images publish.yml pushes, built for this runner's platform only and never pushed. images: diff --git a/AGENTS.md b/AGENTS.md index 4ee2a24..006a8eb 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -37,9 +37,11 @@ docker compose -f docker/compose.yaml up -d worker frankenphp fault-worker hosti docker compose -f docker/compose.yaml run --rm --no-deps php-integration docker compose -f docker/compose.yaml run --rm --no-deps integration docker compose -f docker/compose.yaml --profile test down -docker compose -f docker/compose.yaml --profile mcp --profile mcp-fetch --profile mcp-model --profile php-dev --profile php-frameworks --profile tls config +docker compose -f docker/compose.yaml --profile mcp --profile mcp-fetch --profile mcp-model --profile php-dev --profile py-dev --profile php-frameworks --profile tls config docker compose -f docker/compose.yaml run --rm --no-deps mcp-fetcher # network once: wheels into .cache/ docker compose -f docker/compose.yaml run --rm --no-deps mcp-tests +docker compose -f docker/compose.yaml run --rm --no-deps py-dev-fetcher # network once: mypy and ruff into .cache/ +docker compose -f docker/compose.yaml run --rm --no-deps py-analysis # offline mypy --strict and ruff check docker compose -f docker/compose.yaml run --rm --no-deps composer-install # offline consumer install of the library docker compose -f docker/compose.yaml run --rm --no-deps php-fetcher # network once: PHPStan and the PSR interfaces into .cache/ docker compose -f docker/compose.yaml run --rm --no-deps php-frameworks-fetcher # network once: Laravel and Symfony into .cache/ @@ -50,7 +52,7 @@ docker compose -f docker/compose.yaml run --rm --no-deps laravel-tests # offli docker compose -f docker/compose.yaml run --rm --no-deps symfony-tests # offline: the bundle in a bare Symfony kernel docker compose -f docker/compose.yaml up -d tls-proxy docker compose -f docker/compose.yaml run --rm --no-deps tls-check # the client over HTTPS through Caddy, example CA only -docker compose -f docker/compose.yaml --profile mcp --profile mcp-fetch --profile php-dev --profile php-frameworks --profile tls down +docker compose -f docker/compose.yaml --profile mcp --profile mcp-fetch --profile php-dev --profile py-dev --profile php-frameworks --profile tls down ``` CI runs the PHP steps on each of PHP 8.2, 8.3, 8.4 and 8.5 with `BRIDGE_PHP_IMAGE` set to the matrix's diff --git a/CHANGELOG.md b/CHANGELOG.md index bbaf582..3b2cbf9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,7 @@ All notable changes to this project are documented in this file. The format foll - `GET /readyz`, behind any configured token like `/healthz`: 200 when the worker is not draining and a child is idle with its backend's models loaded, otherwise 503 with a reason (`draining`, `no_idle_child`, `models_loading`). A child now loads its backend's models as it starts instead of on its first job. Compose's MCP workers report healthy only once ready. - Optional structured logging: `BRIDGE_LOG=json` writes one JSON object per line on stderr, best effort (a full queue or the exit can drop lines), a `request` line after each reply (method, route template, status, duration, the token's name), a `job` line when a job ends (job ID, task, owner, terminal status, queue and run times), and `start`, `drain` and `stop` lines. A fixed schema checked value by value keeps inputs, results, headers, query strings, tokens and exception text out; a stderr that blocks or fails never holds a request. `off`, the default, writes nothing, and any other value stops the worker at start naming `BRIDGE_LOG`. - `BRIDGE_OWNER_CAPACITY`, a cap on one token's queued, running and retained jobs (1 to the capacity; unset, the whole capacity, so nothing changes without it). A submission over it is the same 429 `capacity_exceeded` as a full worker, with a `Retry-After` from that token's own first expiry; a repeated idempotency key is still answered, and the total capacity still applies. The `start` log line carries `owner_capacity`. +- Static analysis of the Python code in the required `deterministic` check: `py-analysis` runs `mypy --strict` over the worker and the MCP server and `ruff check` over them, the scripts and the tests, offline, with mypy 2.4.0 and ruff 0.17.2 hash-pinned in `requirements-dev.lock` and configured in `pyproject.toml`; neither is installed at runtime or shipped in the archive. - A full connection budget (32) is answered 503 `service_busy` with `Retry-After: 1` and `Connection: close` instead of a closed socket, and a request line over 8 KiB or a header block over 16 KiB is refused with 431 `headers_too_large`. ### Changed diff --git a/README.md b/README.md index 4126f9e..2b24c8a 100644 --- a/README.md +++ b/README.md @@ -182,6 +182,11 @@ docker compose -f docker/compose.yaml run --rm --no-deps integration docker compose -f docker/compose.yaml run --rm --no-deps mcp-fetcher docker compose -f docker/compose.yaml run --rm --no-deps mcp-tests +# mypy --strict and ruff over the Python code, after acquiring the pinned wheels once (network); +# py-analysis also installs the MCP server's wheels from mcp-fetcher, so it runs after that: +docker compose -f docker/compose.yaml run --rm --no-deps py-dev-fetcher +docker compose -f docker/compose.yaml run --rm --no-deps py-analysis + # PHP syntax validation inside the container: docker compose -f docker/compose.yaml run --rm --no-deps php-tests sh -ec \ 'find src examples tests/php -name "*.php" -print0 | xargs -0 -n1 php -l' @@ -198,7 +203,7 @@ docker compose -f docker/compose.yaml run --rm --no-deps php-analysis Coverage includes invalid contracts, duplicate/non-finite JSON, authentication, oversized bodies, queue deadlines, queued/running cancellation, hard process crashes, retention, malformed/oversized HTTP responses, redirect refusal, and reused-worker request isolation. Fault tasks are enabled only in the test service with `BRIDGE_TEST_TASKS=1`. -CI runs deterministic tests for pushes to main and for pull requests across PHP 8.2–8.5, treating PHP warnings and deprecations as test failures. Every job runs on GitHub's hosted runners (`ubuntu-latest`), because the repository is public: GitHub does not meter Actions minutes for public repositories, and a fork's pull request runs its own code on a disposable runner of GitHub's, not on a machine of the maintainer's; see [the runner notes](docs/testing.md#ci-runners) and [security](docs/security.md). The stable `deterministic` check succeeds only if every PHP matrix job and the MCP server job, which runs the offline MCP suite, the consumer install and PHPStan, succeed. A separate workflow performs model acquisition and the real-model smoke tests on the default PHP 8.5 image: weekly, on pull requests that touch the model path, and on demand. It is not a required check, it does not use paid APIs or GPUs, and it also prints an informational latency table: its jobs must succeed, but its times never gate the run. +CI runs deterministic tests for pushes to main and for pull requests across PHP 8.2–8.5, treating PHP warnings and deprecations as test failures. Every job runs on GitHub's hosted runners (`ubuntu-latest`), because the repository is public: GitHub does not meter Actions minutes for public repositories, and a fork's pull request runs its own code on a disposable runner of GitHub's, not on a machine of the maintainer's; see [the runner notes](docs/testing.md#ci-runners) and [security](docs/security.md). The stable `deterministic` check succeeds only if every PHP matrix job and the MCP server job, which runs the offline MCP suite, mypy and ruff, the consumer install and PHPStan, succeed. A separate workflow performs model acquisition and the real-model smoke tests on the default PHP 8.5 image: weekly, on pull requests that touch the model path, and on demand. It is not a required check, it does not use paid APIs or GPUs, and it also prints an informational latency table: its jobs must succeed, but its times never gate the run. ## Important limits diff --git a/docker/compose.yaml b/docker/compose.yaml index 97a1966..d5d03e1 100644 --- a/docker/compose.yaml +++ b/docker/compose.yaml @@ -618,6 +618,48 @@ services: - -ec - 'python -m pip install -q --require-hashes --no-compile --no-index --find-links=/cache/wheels --only-binary=:all: --target /deps -r requirements-mcp.lock && python -m unittest discover -s tests/mcp -v' + # The pinned mypy and ruff wheels into .cache: the only static-analysis step with network. + py-dev-fetcher: + <<: *python + profiles: [py-dev] + network_mode: bridge + volumes: + - ..:/work:ro + - ../.cache:/cache + environment: + PYTHONDONTWRITEBYTECODE: "1" + PIP_DISABLE_PIP_VERSION_CHECK: "1" + PIP_NO_CACHE_DIR: "1" + command: + - sh + - -ec + - 'python -m pip download --require-hashes --only-binary=:all: --dest /cache/wheels -r requirements-dev.lock' + + # mypy --strict and ruff, offline. The MCP server's wheels (from mcp-fetcher) are installed beside the tools so + # mypy reads the SDK's own types; both tools keep their caches in the tmpfs, since the checkout is read-only. + py-analysis: + <<: *python + profiles: [py-dev] + network_mode: none + mem_limit: 768m + tmpfs: + - /tmp:rw,noexec,nosuid,size=128m,uid=65532,gid=65532 + - /deps:rw,exec,nosuid,size=384m,uid=65532,gid=65532 + volumes: + - ..:/work:ro + - ../.cache:/cache:ro + environment: + PYTHONDONTWRITEBYTECODE: "1" + PYTHONPATH: /deps/tools:/deps/mcp + MYPY_CACHE_DIR: /tmp/mypy + RUFF_CACHE_DIR: /tmp/ruff + PIP_DISABLE_PIP_VERSION_CHECK: "1" + PIP_NO_CACHE_DIR: "1" + command: + - sh + - -ec + - 'python -m pip install -q --require-hashes --no-compile --no-index --find-links=/cache/wheels --only-binary=:all: --target /deps/tools -r requirements-dev.lock && python -m pip install -q --require-hashes --no-compile --no-index --find-links=/cache/wheels --only-binary=:all: --target /deps/mcp -r requirements-mcp.lock && python -m ruff check worker mcp_server scripts tests && python -m mypy' + networks: isolated: internal: true diff --git a/docs/testing.md b/docs/testing.md index e65bad9..e5d2143 100644 --- a/docs/testing.md +++ b/docs/testing.md @@ -29,6 +29,18 @@ docker run --rm --user "$(id -u):$(id -g)" -v "$PWD:/work:ro" -v "$PWD/tests/php 'mkdir /tmp/p && cp /out/composer.json /out/composer.lock /tmp/p && cd /tmp/p && composer update cristiangirlea/php-python-ai-bridge --no-install --no-scripts --no-plugins && cp composer.lock /out/' ``` +`requirements-dev.lock` pins the wheels `py-analysis` installs for the versions in `requirements-dev.txt`, one hash per wheel for CPython 3.12 on Linux x86_64. To move mypy or ruff to a newer release, change `requirements-dev.txt`, then resolve and hash the wheel set again with network in a throwaway container from the repository root, review the diff, and rerun `py-dev-fetcher` and `py-analysis`: + +```sh +docker run --rm -v "$PWD:/work" -w /work python:3.12-bookworm@sha256:581429e3df12d76e6af4be5ab7d0e7fc2013eb57dc23d2de691411c8efdbb970 sh -ec ' +python -m pip download -q --disable-pip-version-check --only-binary=:all: --dest /tmp/w -r requirements-dev.txt +{ echo "# Resolved wheel set for CPython 3.12, Linux x86_64. No source distributions." + for f in /tmp/w/*.whl; do n=$(basename "$f" | cut -d- -f1 | tr _ - | tr A-Z a-z); v=$(basename "$f" | cut -d- -f2) + echo "$n==$v --hash=sha256:$(sha256sum "$f" | cut -d" " -f1)"; done; } > requirements-dev.lock' +``` + +A package both locks pin (today `typing-extensions`) must have the same version in each, since `py-analysis` installs both. + To select another PHP version, set `BRIDGE_PHP_IMAGE` to its exact pinned image from [the CI matrix](../.github/workflows/tests.yml), then run the same Compose commands. For example: ```sh @@ -50,7 +62,7 @@ CI runs on GitHub's hosted runners, `ubuntu-latest` for every job, because the r Every job starts on a fresh runner with an empty image store and pulls `python:3.12-bookworm` and the FrankenPHP images, and the MCP job `composer:2`, from Docker Hub, which limits unauthenticated pulls per IP; GitHub's runners share address ranges, so a pull that fails with a rate-limit error is a rerun, not a regression. The model workflow keeps both model workers up, each allowed 1536 MiB with its wheels in tmpfs, beside FrankenPHP, the MCP server, the demo worker and one client, whose limits add up to about 5 GiB; a Docker daemon that reproduces that workflow locally needs about 6 GiB. -CI checks the actual PHP minor before testing so a misconfigured image cannot silently test the default version four times. The stable required check `deterministic` depends on all four matrix jobs and on the MCP server job, which also runs the consumer install, PHPStan, the PSR-18 transport checks, the Laravel and Symfony tests and the TLS example, and on the `images` job, which builds both worker images for linux/amd64 and starts each, and rejects any of them failed, cancelled or skipped. `publish.yml` pushes those images for amd64 and arm64 on a `v*` tag only. PHP contract/HTTP tests promote warnings and deprecations to exceptions. Dependabot proposes weekly updates through `.github/dependabot.yml` for three ecosystems: the GitHub Actions pins, the pip wheel sets (`requirements-model.*`, `requirements-mcp.*`) and `composer.lock`. +CI checks the actual PHP minor before testing so a misconfigured image cannot silently test the default version four times. The stable required check `deterministic` depends on all four matrix jobs and on the MCP server job, which also runs mypy and ruff, the consumer install, PHPStan, the PSR-18 transport checks, the Laravel and Symfony tests and the TLS example, and on the `images` job, which builds both worker images for linux/amd64 and starts each, and rejects any of them failed, cancelled or skipped. `publish.yml` pushes those images for amd64 and arm64 on a `v*` tag only. PHP contract/HTTP tests promote warnings and deprecations to exceptions. Dependabot proposes weekly updates through `.github/dependabot.yml` for three ecosystems: the GitHub Actions pins, the pip wheel sets (`requirements-model.*`, `requirements-mcp.*`) and `composer.lock`. ## What the tests establish @@ -66,6 +78,7 @@ PHP branch support was verified locally on 2026-09-17 on PHP 8.2.33, 8.3.33, 8.4 | PHP syntax | All 25 PHP source/example/test files | | Consumer install | `composer validate --strict` on the library, then a throwaway consumer project in tmpfs, with Packagist disabled and the platform pinned to PHP 8.2.0, installs its `git archive` zip from an `artifact` repository (or from a path repository in a git worktree whose git directory is outside the mount; the archive is built from the committed `HEAD`, so an uncommitted change is not what it tests), checks that `vendor/` holds `src/Client.php` and `docs/` and not `worker/`, `tests/` or `docker/`, that the library brought no other vendor into `vendor/` (so no framework), and autoloads `PhpAiBridge\Client`; runs once, in the MCP server job, with no network | | PHPStan | Level max over `src`, `examples` and `tests/php` with `phpstan.neon`, PHPStan 2.3.1 from `composer.lock` installed into `.cache/php-vendor` by `php-fetcher` and run offline by `php-analysis`, with the frameworks' autoloader from `.cache/php-frameworks-vendor` as its bootstrap file so the provider, the bundle and their tests are analysed against Laravel and Symfony; `phpstan-baseline.neon` holds 9 findings in examples and tests, each with its reason, and none in `src`; runs once, in the MCP server job | +| Python static analysis | `ruff check` over `worker`, `mcp_server`, `scripts` and `tests` with pycodestyle's errors and warnings, pyflakes, bugbear and the import order (`E`, `W`, `F`, `B`, `I`, but not `E501`, line length, or `B905`, `zip()` without `strict=`), then `mypy --strict` over `worker` and `mcp_server` with the MCP SDK's own types; mypy 2.4.0 and ruff 0.17.2 from `requirements-dev.lock`, downloaded by `py-dev-fetcher` and installed offline with the MCP server's wheels into a tmpfs by `py-analysis`, configured in `pyproject.toml`; `numpy`, `onnxruntime` and `tokenizers` are not installed there and are typed as unknown, and three `type: ignore` comments remain, each with its reason; runs once, in the MCP server job | | Laravel and Symfony | `php-frameworks-fetcher` installs what `tests/php/frameworks/composer.lock` pins (Laravel 13.38.0, Symfony 8.0.15, resolved for PHP 8.4 and later) into `.cache/php-frameworks-vendor`, with the library linked from the checkout by a path repository, as Composer installs a package; then, offline, against the real `worker`: `laravel-tests`, 94 checks, boots an `Illuminate\Foundation\Application` with no skeleton whose provider comes from package discovery, publishes the configuration with `vendor:publish`, resolves one client twice, reranks, boots again from a configuration cache with the environment changed, and refuses a missing, empty or malformed URL, a missing, empty or short token and a timeout of 0, 300001 or not a number, naming the setting and never quoting the token; `symfony-tests`, 96 checks, boots a bare `Kernel` with only the bundle and debug off, resolves the client and an autowired service, reranks, loads the dumped container in a second kernel, and refuses the same cases, plus an empty URL and a zero timeout written in the configuration at build time; once, on PHP 8.5, in the MCP server job; no full application, FrameworkBundle, Octane, Messenger or queue worker runs | | TLS proxy | `tls-check` runs `tests/php/tls.php` against `tls-proxy`, Caddy with `tls internal` in front of the real `worker`: with the system CA store emptied by a tmpfs, so the client trusts only the example root: 7 checks with `curl.cainfo` set to that root (that setting itself, the empty system store, a rerank over HTTPS with the certificate verified, 401 for a wrong token, the worker's own 401 for a request with no `Authorization` header, Caddy's 400 for plain HTTP on the TLS port, no listener on port 80) and 3 with no CA configured (no example CA set, a TLS handshake without verification succeeds, the client fails as `transport_error`); once, on PHP 8.5, in the MCP server job | | Optional ONNX | Through both the PHP CLI and FrankenPHP HTTP: 2 rerank cases, a 70-document rerank across three batches with its relevant document in the second, a parity check that three documents score within 1e-3 alone and inside that set, 1 embedding and 3 redaction cases: a name, an email address and a place; a sentence whose lower-case person, place and organisation only the uncased model finds, which also checks its label order; and clean text, which must produce no spans; and the ONNX-backed MCP server end to end, every tool through compose's `mcp-model` service over stdio, including a search over an index the embedding model built, with a query worded unlike the note it must find, and the calls for every question in the semantic evaluation, whose answers the models must all reach; smoke checks, not quality datasets | @@ -87,6 +100,7 @@ The original prototype was not entirely developed with TDD. Review corrections u - Ranking quality is measured on one public dataset, ArguAna, a task unlike the reranker's training; none on short-query search and none in another language. NER quality on real data is one public set, WNUT-17's English social media text, with four of its six types merged into three labels; no other domain or language, and no GPU benchmark. The invented-sample NER figures below come from 48 sentences by the same author as the development set used to tune the redact task. - No repeated or semantic LLM evaluation. The one recorded run below is a single session per question per model on the demo worker, so it shows that the tools can be used, not a success rate or retrieval quality. - No real certificate and no mutual TLS: the TLS example signs with Caddy's local CA, runs on one Docker network rather than two hosts, and checks no client certificate. +- No type checking of `scripts/` and `tests/`: mypy covers the worker and the MCP server only, and ruff all four. - No verified native Windows/macOS service execution or ARM runtime; Docker tests target Linux amd64/Python 3.12. - No durable delivery, replay or automatic retry, and no multi-instance result routing. With `BRIDGE_STATE_DIR` set the job table (never inputs) survives a restart, so retained results and idempotency keys are answered and unfinished jobs are reported as `service_restarted`; without it, jobs and results live in memory. diff --git a/mcp_server/bridge_mcp/__main__.py b/mcp_server/bridge_mcp/__main__.py index 397f1a7..7285702 100644 --- a/mcp_server/bridge_mcp/__main__.py +++ b/mcp_server/bridge_mcp/__main__.py @@ -4,12 +4,13 @@ import sys import time from pathlib import Path +from typing import Any from .protocol import Bridge, BridgeError from .tools import build_server -def wait_for_worker(bridge: Bridge, window_s: float) -> dict: +def wait_for_worker(bridge: Bridge, window_s: float) -> dict[str, Any]: """Compose starts the worker beside this server with no readiness gate, so a refused connection is waited out for a bounded window; a rejected token or a bad contract is not.""" deadline = time.monotonic() + window_s diff --git a/mcp_server/bridge_mcp/index.py b/mcp_server/bridge_mcp/index.py index 003a564..94dfcab 100644 --- a/mcp_server/bridge_mcp/index.py +++ b/mcp_server/bridge_mcp/index.py @@ -24,8 +24,10 @@ import threading import time from collections import OrderedDict +from collections.abc import Callable, Iterable, Iterator, Sequence from dataclasses import dataclass, field from pathlib import Path +from typing import IO, Any from .protocol import Bridge, BridgeError @@ -63,14 +65,14 @@ class LoadedIndex: backend: str dimensions: int count: int - chunking: dict - sources: list - chunks: list - vectors: array.array + chunking: dict[str, Any] + sources: list[str] + chunks: list[dict[str, Any]] + vectors: array.array[float] directory: Path vectors_path: Path = field(repr=False) # (device, inode, mtime, size) of the sidecar and vectors files exactly as read. - stamps: tuple = field(repr=False, default=()) + stamps: tuple[tuple[int, int, int, int], ...] = field(repr=False, default=()) @property def size(self) -> int: @@ -84,7 +86,7 @@ def _skip_space(text: str, i: int) -> int: return word.start() if word else len(text) -def _last_space(text: str, low: int, high: int): +def _last_space(text: str, low: int, high: int) -> int | None: """The index of the last whitespace character in text[low:high + 1], or None. Unicode whitespace counts, as it does everywhere else in the chunker.""" last = None @@ -104,7 +106,7 @@ def _next_word(text: str, i: int) -> int: return _skip_space(text, i) -def chunk_text(text: str, chars: int = 1000, overlap: int = 200) -> list: +def chunk_text(text: str, chars: int = 1000, overlap: int = 200) -> list[tuple[int, int]]: """Character windows of at most `chars`, cut at whitespace where one exists in the second half of the window, each starting on a word and overlapping the previous one by about `overlap` characters. @@ -144,8 +146,8 @@ def _check(condition: bool, message: str) -> None: raise InvalidIndex(message) -def write_index(out: Path, *, model: str, backend: str, dimensions: int, chunking: dict, sources: list, - chunks: list, rows) -> dict: +def write_index(out: Path, *, model: str, backend: str, dimensions: int, chunking: dict[str, Any], sources: list[str], + chunks: list[dict[str, Any]], rows: Iterable[Iterable[float]]) -> dict[str, Any]: """Stream `rows` into vectors.f32, then write index.json. Nothing is replaced unless both were written; the two renames are separate, so a reader between them sees a checksum mismatch and must retry.""" out = Path(out) @@ -241,7 +243,7 @@ def load_index(sidecar: Path, within: Path | None = None) -> LoadedIndex: sidecar.parent, vectors_path, (sidecar_stamp, vectors_stamp)) -def search(loaded: LoadedIndex, query_vector, n: int) -> list: +def search(loaded: LoadedIndex, query_vector: Iterable[float], n: int) -> list[tuple[int, float]]: """The n chunks with the highest dot product, as (chunk, score); ties keep the lower chunk number.""" # A memoryview slices rows without copying them. width, vectors = loaded.dimensions, memoryview(loaded.vectors) @@ -250,16 +252,16 @@ def search(loaded: LoadedIndex, query_vector, n: int) -> list: return [(-negative, score) for score, negative in heapq.nlargest(n, scores)] -def _identity(status) -> tuple: +def _identity(status: os.stat_result) -> tuple[int, int, int, int]: # Every rebuild replaces both files by rename, so the inode changes even when size and time do not. return status.st_dev, status.st_ino, status.st_mtime_ns, status.st_size -def _fstamp(handle) -> tuple: +def _fstamp(handle: IO[bytes]) -> tuple[int, int, int, int]: return _identity(os.fstat(handle.fileno())) -def _stamp(path: Path) -> tuple: +def _stamp(path: Path) -> tuple[int, int, int, int]: return _identity(path.stat()) @@ -269,7 +271,7 @@ class IndexCache: def __init__(self, limit_bytes: int = CACHE_BYTES): self.limit_bytes = limit_bytes - self._items = OrderedDict() + self._items: OrderedDict[tuple[Path, Path | None], LoadedIndex] = OrderedDict() self._lock = threading.Lock() def get(self, sidecar: Path, within: Path | None = None) -> LoadedIndex: @@ -307,7 +309,7 @@ def _pause(seconds: float) -> None: time.sleep(seconds) -def collect(inputs: list, suffixes=DEFAULT_SUFFIXES) -> list: +def collect(inputs: Sequence[str | Path], suffixes: tuple[str, ...] = DEFAULT_SUFFIXES) -> list[Path]: """Every file named, and every file with one of `suffixes` under every directory named, sorted.""" files = [] for item in inputs: @@ -324,8 +326,9 @@ def collect(inputs: list, suffixes=DEFAULT_SUFFIXES) -> list: return unique -def _batches(texts: list): - batch, size = [], 0 +def _batches(texts: list[str]) -> Iterator[list[int]]: + batch: list[int] = [] + size = 0 for number, text in enumerate(texts): encoded = len(json.dumps(text, ensure_ascii=False).encode("utf-8")) + 1 if batch and (len(batch) == EMBED_BATCH or size + encoded > BATCH_BYTES): @@ -337,7 +340,7 @@ def _batches(texts: list): yield batch -def _embed(bridge, texts: list, timeout_s: float, identity: dict) -> list: +def _embed(bridge: Bridge, texts: list[str], timeout_s: float, identity: dict[str, Any]) -> list[list[float]]: for attempt in range(CAPACITY_RETRIES + 1): try: result = bridge.run("embed", {"texts": texts}, timeout_s) @@ -362,15 +365,17 @@ def _embed(bridge, texts: list, timeout_s: float, identity: dict) -> list: return vectors -def build(bridge, inputs: list, out: Path, chars: int = 1000, overlap: int = 200, - suffixes=DEFAULT_SUFFIXES, timeout_s: float = 60.0, log=None) -> dict: +def build(bridge: Bridge, inputs: Sequence[str | Path], out: Path, chars: int = 1000, overlap: int = 200, + suffixes: tuple[str, ...] = DEFAULT_SUFFIXES, timeout_s: float = 60.0, + log: Callable[[str], None] | None = None) -> dict[str, Any]: """Chunk every input file, embed the chunks through the worker, and write the index into `out`.""" out = Path(out) if out.exists() and not out.is_dir(): raise BuildError(f"--out is an existing file: {out}") files = collect(inputs, suffixes) backend = bridge.health()["backend"] - sources, chunks = [], [] + sources: list[str] = [] + chunks: list[dict[str, Any]] = [] for number, path in enumerate(files): if path.stat().st_size > MAX_SOURCE_BYTES: raise BuildError(f"{path} is larger than {MAX_SOURCE_BYTES // (1024 * 1024)} MiB") @@ -388,9 +393,9 @@ def build(bridge, inputs: list, out: Path, chars: int = 1000, overlap: int = 200 if len(chunks) > MAX_INDEX_CHUNKS: raise BuildError(f"{len(chunks)} chunks exceed the {MAX_INDEX_CHUNKS}-chunk limit; index fewer files") texts = [chunk["text"] for chunk in chunks] - identity = {} + identity: dict[str, Any] = {} - def rows(): + def rows() -> Iterator[list[float]]: done = 0 for batch in _batches(texts): yield from _embed(bridge, [texts[number] for number in batch], timeout_s, identity) @@ -407,12 +412,12 @@ def rows(): return meta -def _prepend(first, rest): +def _prepend(first: list[float], rest: Iterator[list[float]]) -> Iterator[list[float]]: yield first yield from rest -def main(argv=None) -> int: +def main(argv: Sequence[str] | None = None) -> int: parser = argparse.ArgumentParser(prog="python -m bridge_mcp.index") commands = parser.add_subparsers(dest="command", required=True) command = commands.add_parser("build", help="build an index from UTF-8 text files") diff --git a/mcp_server/bridge_mcp/protocol.py b/mcp_server/bridge_mcp/protocol.py index 30a303c..4533614 100644 --- a/mcp_server/bridge_mcp/protocol.py +++ b/mcp_server/bridge_mcp/protocol.py @@ -4,6 +4,8 @@ import json import re import time +from collections.abc import Callable +from typing import Any from urllib.parse import urlsplit MAX_BYTES = 262144 @@ -39,13 +41,13 @@ def __init__(self, base_url: str, token: str, request_timeout_s: float = 5.0): self._token = token self._timeout = request_timeout_s - def health(self) -> dict: + def health(self) -> dict[str, Any]: body = self._request("GET", "/healthz") if body.get("protocol") != 1 or not isinstance(body.get("backend"), str): raise BridgeError("Unexpected health response", "invalid_response") return body - def submit(self, task: str, payload: dict, timeout_ms: int) -> dict: + def submit(self, task: str, payload: dict[str, Any], timeout_ms: int) -> dict[str, Any]: if not 100 <= timeout_ms <= 300000: raise ValueError("timeout_ms must be between 100 and 300000") job = self._job(self._request("POST", "/v1/jobs", {"task": task, "input": payload, "timeout_ms": timeout_ms})) @@ -53,13 +55,14 @@ def submit(self, task: str, payload: dict, timeout_ms: int) -> dict: raise BridgeError("Submission response does not match the request", "invalid_response") return job - def get(self, job_id: str) -> dict: + def get(self, job_id: str) -> dict[str, Any]: return self._job(self._request("GET", "/v1/jobs/" + self._id(job_id)), job_id) - def cancel(self, job_id: str) -> dict: + def cancel(self, job_id: str) -> dict[str, Any]: return self._job(self._request("POST", "/v1/jobs/" + self._id(job_id) + "/cancel", {}), job_id) - def run(self, task: str, payload: dict, deadline_s: float, on_progress=None, poll_s: float = 0.1) -> dict: + def run(self, task: str, payload: dict[str, Any], deadline_s: float, + on_progress: Callable[[int, int], object] | None = None, poll_s: float = 0.1) -> dict[str, Any]: """Submit, poll to a terminal state and return the result. A local deadline cancels the remote job, unlike the PHP client's wait(): a web request can @@ -91,7 +94,8 @@ def run(self, task: str, payload: dict, deadline_s: float, on_progress=None, pol pass raise if job["status"] == "succeeded": - return job["result"] + result: dict[str, Any] = job["result"] + return result code = (job.get("error") or {}).get("code", job["status"]) raise BridgeError(f"Job {job['status']}: {code}", code, job_id=job["id"]) @@ -102,7 +106,7 @@ def _id(job_id: str) -> str: return job_id @staticmethod - def _job(body: dict, expected_id: str | None = None) -> dict: + def _job(body: dict[str, Any], expected_id: str | None = None) -> dict[str, Any]: keys = ("id", "task", "status", "result", "error", "progress", "timeout_ms") if any(key not in body for key in keys): raise BridgeError("Incomplete job response", "invalid_response") @@ -122,7 +126,7 @@ def _job(body: dict, expected_id: str | None = None) -> dict: raise BridgeError("Job ID does not match the request", "invalid_response") return body - def _request(self, method: str, path: str, body: dict | None = None) -> dict: + def _request(self, method: str, path: str, body: dict[str, Any] | None = None) -> dict[str, Any]: data = None if body is not None: data = json.dumps(body, ensure_ascii=False, allow_nan=False, separators=(",", ":")).encode("utf-8") diff --git a/mcp_server/bridge_mcp/tools.py b/mcp_server/bridge_mcp/tools.py index a2c4a0f..918bace 100644 --- a/mcp_server/bridge_mcp/tools.py +++ b/mcp_server/bridge_mcp/tools.py @@ -1,13 +1,14 @@ """One MCP tool per allowlisted task. The tool list is a literal in this file: nothing is read from the worker, and the test-only tasks cannot appear whatever BRIDGE_TEST_TASKS says.""" +import functools import json import math import random import re import time from pathlib import Path -from typing import Annotated, Literal, TypedDict +from typing import Annotated, Any, Literal, TypedDict, cast import anyio from mcp.server import MCPServer @@ -184,20 +185,21 @@ def _explain(error: BridgeError) -> str: return hints.get(error.code, f"The worker reported {error.code}; internal details are not exposed.") -def _number(value) -> bool: - return type(value) in (int, float) and math.isfinite(value) +def _number(value: object) -> bool: + return type(value) in (int, float) and math.isfinite(cast(float, value)) -def _name(value) -> bool: +def _name(value: object) -> bool: return isinstance(value, str) and value != "" -def _rankings(result: dict, count: int, top_k: int) -> list: +def _rankings(result: dict[str, Any], count: int, top_k: int) -> list[dict[str, Any]]: # The same rules as RerankResult in the PHP client: exact count, unique indexes, stable descending order. rankings = result.get("rankings") if not _name(result.get("model")) or not isinstance(rankings, list) or len(rankings) != min(top_k, count): raise ToolError("The worker returned an invalid ranking.") - seen, previous = set(), None + seen: set[int] = set() + previous: tuple[int, float] | None = None for item in rankings: if (not isinstance(item, dict) or type(item.get("index")) is not int or not 0 <= item["index"] < count or item["index"] in seen or not _number(item.get("score")) @@ -206,10 +208,10 @@ def _rankings(result: dict, count: int, top_k: int) -> list: raise ToolError("The worker returned an invalid ranking.") seen.add(item["index"]) previous = (item["index"], item["score"]) - return rankings + return cast(list[dict[str, Any]], rankings) -def _vectors(result: dict, count: int) -> list: +def _vectors(result: dict[str, Any], count: int) -> list[list[float]]: # The same rules as EmbedResult: one vector of the declared dimension per text, finite, unit length. vectors, dimensions = result.get("vectors"), result.get("dimensions") if (not _name(result.get("model")) or type(dimensions) is not int or dimensions < 1 @@ -219,10 +221,10 @@ def _vectors(result: dict, count: int) -> list: if (not isinstance(vector, list) or len(vector) != dimensions or not all(_number(value) for value in vector) or abs(math.sqrt(sum(value * value for value in vector)) - 1.0) > 1e-3): raise ToolError("The worker returned invalid embeddings.") - return vectors + return cast(list[list[float]], vectors) -def _spans(result: dict, text: str) -> list: +def _spans(result: dict[str, Any], text: str) -> list[Span]: # The same rules as RedactResult: spans sorted, disjoint and inside the original text. spans = result.get("spans") if (not _name(result.get("model")) or not isinstance(result.get("text"), str) or not result["text"] @@ -237,10 +239,10 @@ def _spans(result: dict, text: str) -> list: or not _number(span.get("score")) or not 0 <= span["score"] <= 1): raise ToolError("The worker returned an invalid redaction.") cursor = span["end"] - return spans + return cast(list[Span], spans) -def _check_texts(texts: list, singular: str, limit: int, extra: int = 0) -> None: +def _check_texts(texts: list[str], singular: str, limit: int, extra: int = 0) -> None: if not 1 <= len(texts) <= limit: raise ToolError(f"Give between 1 and {limit} {singular}s.") total = extra @@ -298,7 +300,7 @@ def _documents_from(root: Path | None, path_text: str) -> list[str]: return documents -def _index_from(root: Path | None, path_text: str, cache) -> index_files.LoadedIndex: +def _index_from(root: Path | None, path_text: str, cache: index_files.IndexCache) -> index_files.LoadedIndex: resolved = _confine(root, path_text, "index_path") sidecar = resolved try: @@ -368,7 +370,7 @@ def build_server(bridge: Bridge, backend: str, root: Path | None, timeout_ms: in mcp = MCPServer("bridge_mcp", title="PHP-Python AI bridge", version="0.1", instructions=INSTRUCTIONS.format(backend=backend, **described)) - async def run(task: str, payload: dict, ctx: Context, deadline: float | None = None) -> dict: + async def run(task: str, payload: dict[str, Any], ctx: Context, deadline: float | None = None) -> dict[str, Any]: """Run one job within the call's time budget; a tool that runs several jobs passes one deadline to all.""" def on_progress(completed: int, total: int) -> None: # The SDK sends nothing when the host gave no progress token, so this is unconditional. @@ -382,7 +384,7 @@ def on_progress(completed: int, total: int) -> None: remaining = deadline - time.monotonic() try: return await anyio.to_thread.run_sync( - lambda remaining=remaining: bridge.run(task, payload, remaining, on_progress, 0.1)) + functools.partial(bridge.run, task, payload, remaining, on_progress, 0.1)) except BridgeError as error: if error.code != "capacity_exceeded": raise ToolError(_explain(error)) from None @@ -443,7 +445,11 @@ async def bridge_rerank( raise ToolError("query is empty.") if (documents is None) == (documents_path is None): raise ToolError("Give exactly one of documents or documents_path.") - chosen = documents if documents is not None else _documents_from(root, documents_path) + if documents is None: + assert documents_path is not None # exactly one of the two, as checked above + chosen = _documents_from(root, documents_path) + else: + chosen = documents _check_texts(chosen, "document", BY_VALUE_DOCUMENTS if documents is not None else BY_REFERENCE_DOCUMENTS, extra=len(query)) result = await run("rerank", {"query": query, "documents": chosen, "top_k": min(top_k, len(chosen))}, ctx) @@ -469,7 +475,7 @@ async def bridge_embed_similarity( result = await run("embed", {"texts": texts}, ctx) vectors = _vectors(result, len(texts)) matrix = [[round(sum(x * y for x, y in zip(a, b)), 4) for b in vectors] for a in vectors] - pairs = sorted(({"a": a, "b": b, "similarity": matrix[a][b]} + pairs = sorted((Pair(a=a, b=b, similarity=matrix[a][b]) for a in range(len(texts)) for b in range(a + 1, len(texts))), key=lambda pair: (-pair["similarity"], pair["a"], pair["b"]))[:PAIRS] return {"model": result["model"], "matrix": matrix, "pairs": pairs} @@ -486,7 +492,7 @@ async def bridge_redact( text: Annotated[str, Field(description="The text to mask.", min_length=1, max_length=MAX_TEXT_CHARACTERS)], ctx: Context, entities: Annotated[list[Literal["PER", "ORG", "LOC"]], Field( - description="Entity kinds the model pass masks; rules always run.", max_length=3)] = ("PER",), + description="Entity kinds the model pass masks; rules always run.", max_length=3)] = ("PER",), # type: ignore[assignment] # a tuple, since a list default would be shared by every call min_score: Annotated[float, Field(description="Minimum model confidence to mask.", ge=0, le=1)] = 0.85, ) -> Redaction: if not text.strip(): @@ -518,6 +524,7 @@ async def bridge_search( ) -> Search: _check_texts([query], "query", 1) loaded = await anyio.to_thread.run_sync(lambda: _index_from(root, index_path, cache)) + assert root is not None # _index_from refuses every path without one # Vectors from different models are not comparable, so the gate comes before any job is spent. if loaded.model != names["embed"]: raise ToolError(f"The index at {index_path!r} was built with {loaded.model!r}; this worker embeds with " @@ -531,7 +538,8 @@ async def bridge_search( f"{loaded.model!r}. The worker changed since start-up; check bridge_health.") wanted = min(loaded.count, max(CANDIDATE_FACTOR * top_k, MIN_CANDIDATES)) candidates = await anyio.to_thread.run_sync(lambda: index_files.search(loaded, query_vector, wanted)) - rerank_model, ordered = None, [(number, similarity, None) for number, similarity in candidates] + rerank_model = None + ordered: list[tuple[int, float, float | None]] = [(number, similarity, None) for number, similarity in candidates] if rerank: # Keep the query and the candidates inside the worker's character budget and the request's byte # limit; multi-byte text reaches the bytes first. @@ -550,7 +558,7 @@ async def bridge_search( rerank_model = reranked["model"] ordered = [(candidates[item["index"]][0], candidates[item["index"]][1], item["score"]) for item in _rankings(reranked, len(texts), wanted_k)] - results = [] + results: list[SearchHit] = [] for number, similarity, score in ordered[:top_k]: chunk = loaded.chunks[number] results.append({"source": _source(loaded, root, chunk["source"]), "chunk": number, diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..3ea638e --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,27 @@ +# Tool configuration only: this repository ships a PHP library and builds no Python package. + +[tool.mypy] +python_version = "3.12" +strict = true +mypy_path = ["worker", "mcp_server"] +files = ["worker", "mcp_server"] + +# The ONNX backend imports these inside its functions only, and they are not installed where mypy runs. +[[tool.mypy.overrides]] +module = ["numpy", "numpy.*", "onnxruntime", "onnxruntime.*", "tokenizers", "tokenizers.*"] +ignore_missing_imports = true + +[tool.ruff] +target-version = "py312" +line-length = 120 +src = ["worker", "mcp_server", "scripts", "tests", "tests/mcp", "tests/python"] + +[tool.ruff.lint] +# pycodestyle errors and warnings, pyflakes, bugbear and import order. E501 is left out: tests and scripts hold +# long literal data and messages that wrapping would restyle without making anything safer. +select = ["E", "W", "F", "B", "I"] +ignore = ["B905", "E501"] + +[tool.ruff.lint.isort] +# The MCP SDK, not the tests/mcp directory that src lists for the test helpers' imports. +known-third-party = ["mcp"] diff --git a/requirements-dev.lock b/requirements-dev.lock new file mode 100644 index 0000000..d895d32 --- /dev/null +++ b/requirements-dev.lock @@ -0,0 +1,8 @@ +# Resolved wheel set for CPython 3.12, Linux x86_64. No source distributions. +ast-serialize==0.13.1 --hash=sha256:9f99786e0d62082a5b121b3437f3167562133711df22128de1e9886f57f0b226 +librt==0.16.0 --hash=sha256:36e53948e99bbe3ffea257124cfcae1cfb01831555c9a9c903c9f9a72db7fd07 +mypy==2.4.0 --hash=sha256:29eb0b9427a6b11b992e452f6cceb8af724f4dceb47e779d0b35e405e996ea5e +mypy-extensions==1.1.0 --hash=sha256:1be4cccdb0f2482337c4743e60421de3a356cd97508abadd57d47403e94f5505 +pathspec==1.1.1 --hash=sha256:a00ce642f577bf7f473932318056212bc4f8bfdf53128c78bbd5af0b9b20b189 +ruff==0.17.2 --hash=sha256:c60f93fec7142dba5b8f487f94bf2f6231dae24dc11e4c48648347c8b0720004 +typing-extensions==4.16.0 --hash=sha256:481caa481374e813c1b176ada14e97f1f67a4539ce9cfeb3f350d78d6370c2e8 diff --git a/requirements-dev.txt b/requirements-dev.txt new file mode 100644 index 0000000..6f5de5f --- /dev/null +++ b/requirements-dev.txt @@ -0,0 +1,4 @@ +# Static analysis of the worker, the MCP server, the scripts and the tests. Never installed at runtime: the +# py-analysis service installs these wheels offline into a tmpfs, hash-checked against requirements-dev.lock. +mypy==2.4.0 +ruff==0.17.2 diff --git a/scripts/mcp_evaluate.py b/scripts/mcp_evaluate.py index fc395b3..5cfdd36 100644 --- a/scripts/mcp_evaluate.py +++ b/scripts/mcp_evaluate.py @@ -37,6 +37,7 @@ BACKENDS = {"lexical": ("mcp_stdio.sh", "mcp-index"), "onnx": ("mcp_stdio_onnx.sh", "mcp-model-index")} sys.path.insert(0, str(REPO / "worker")) from ai_bridge.tasks import MODEL_NAMES # noqa: E402 (standard library only at import) + # Set by a Claude Code session for its children; a nested `claude` started with them believes it is inside one. NESTING = ("CLAUDECODE", "CLAUDE_CODE_ENTRYPOINT") diff --git a/scripts/ner_measure.py b/scripts/ner_measure.py index 531f6b4..6e85f02 100644 --- a/scripts/ner_measure.py +++ b/scripts/ner_measure.py @@ -23,8 +23,8 @@ from urllib.error import HTTPError sys.path.insert(0, str(Path(__file__).resolve().parent)) -from benchmark import POLL_S, TERMINAL, call, cpu_model, wait_ready # noqa: E402 (the same protocol helpers) import public_data # noqa: E402 +from benchmark import POLL_S, TERMINAL, call, cpu_model, wait_ready # noqa: E402 (the same protocol helpers) REPO = Path(__file__).resolve().parents[1] SAMPLE = REPO / "tests" / "fixtures" / "ner-sample.json" @@ -146,7 +146,7 @@ def finish(base: str, body: dict) -> dict: if error.code != 429: raise if attempt == CAPACITY_RETRIES: - raise SystemExit(f"the worker stayed at capacity for {CAPACITY_RETRIES * CAPACITY_WAIT_S} s") + raise SystemExit(f"the worker stayed at capacity for {CAPACITY_RETRIES * CAPACITY_WAIT_S} s") from None time.sleep(CAPACITY_WAIT_S) while job["status"] not in TERMINAL: time.sleep(POLL_S) diff --git a/scripts/relevance_measure.py b/scripts/relevance_measure.py index a1ce74a..f307dd0 100644 --- a/scripts/relevance_measure.py +++ b/scripts/relevance_measure.py @@ -24,12 +24,13 @@ sys.path.insert(0, str(Path(__file__).resolve().parent)) sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "worker")) +import public_data # noqa: E402 + # The worker's own limits, so that no job is refused for its size. from ai_bridge.server import MAX_BODY # noqa: E402 -from ai_bridge.tasks import MAX_DOCUMENTS, MAX_FIELD_CHARACTERS, MAX_TOTAL_CHARACTERS # noqa: E402 +from ai_bridge.tasks import MAX_DOCUMENTS, MAX_FIELD_CHARACTERS, MAX_TOTAL_CHARACTERS # noqa: E402, F401 (the tests hold CANDIDATES to MAX_DOCUMENTS) from benchmark import cpu_model, wait_ready # noqa: E402 from ner_measure import finish # noqa: E402 (submits a job and waits out a full worker) -import public_data # noqa: E402 DATASET = "arguana" QUERIES, CANDIDATES, K = 200, 100, 10 @@ -168,7 +169,7 @@ def main() -> None: lists = [candidates(index, ids, query_id, text, CANDIDATES) for query_id, text in chosen] bm25_seconds = time.monotonic() - started reranked, models, shortened, seconds = [], set(), 0, [] - for number, ((query_id, text), found) in enumerate(zip(chosen, lists), start=1): + for number, ((_query_id, text), found) in enumerate(zip(chosen, lists), start=1): if not found: reranked.append([]) continue diff --git a/tests/mcp/test_index.py b/tests/mcp/test_index.py index 7bca78a..196d831 100644 --- a/tests/mcp/test_index.py +++ b/tests/mcp/test_index.py @@ -60,7 +60,7 @@ def test_unicode_whitespace_is_a_cut_point(self): text = "\u3000".join(f"word{i}" for i in range(400)) spans = index.chunk_text(text, 1000, 200) self.assertGreater(len(spans), 2) - for start, end in spans[:-1]: + for _start, end in spans[:-1]: self.assertTrue(text[end].isspace(), "an ideographic space is a word boundary") def test_parameters_are_bounded(self): @@ -116,7 +116,9 @@ def edit(self, sidecar, **fields): sidecar.write_text(json.dumps(meta), encoding="utf-8") def test_every_corruption_is_refused(self): - vectors = lambda sidecar: sidecar.parent / index.VECTORS + def vectors(sidecar): + return sidecar.parent / index.VECTORS + self.corrupt(lambda s: self.edit(s, format="bridge-index/2")) self.corrupt(lambda s: self.edit(s, model="")) self.corrupt(lambda s: self.edit(s, dimensions=4)) diff --git a/tests/mcp_model_probe.py b/tests/mcp_model_probe.py index 2f50239..534b139 100644 --- a/tests/mcp_model_probe.py +++ b/tests/mcp_model_probe.py @@ -22,8 +22,8 @@ sys.path.insert(0, str(REPO / "tests" / "mcp")) sys.path.insert(0, str(REPO / "scripts")) from ai_bridge.tasks import MODEL_NAMES # noqa: E402 (standard library only at import) -from model_cases import REDACT_TEXT, RERANK_CASES, SIMILAR_TEXTS # noqa: E402 from mcp_evaluate import load_pairs # noqa: E402 (the runner's rule for which answer a backend expects) +from model_cases import REDACT_TEXT, RERANK_CASES, SIMILAR_TEXTS # noqa: E402 from semantic_reach import SEMANTIC_REACH # noqa: E402 SEMANTIC = REPO / "tests" / "mcp" / "evaluation-semantic.xml" diff --git a/tests/model_integration.py b/tests/model_integration.py index 01456fc..fa1f002 100644 --- a/tests/model_integration.py +++ b/tests/model_integration.py @@ -3,7 +3,6 @@ from integration import finished, request, wait_ready from model_cases import REDACT_TEXT, RERANK_CASES, SIMILAR_TEXTS, UNCASED_MASKED, UNCASED_TEXT - FRANCE = ["Bananas are yellow fruit.", "Paris is the capital of France.", "A car has four wheels."] @@ -51,7 +50,10 @@ def main(): assert status == 202 result = finished(job["id"])["result"] assert result["model"] == "sentence-transformers/all-MiniLM-L6-v2" and result["dimensions"] == 384 - dot = lambda a, b: sum(x * y for x, y in zip(a, b)) + + def dot(a, b): + return sum(x * y for x, y in zip(a, b)) + assert dot(result["vectors"][0], result["vectors"][1]) > dot(result["vectors"][0], result["vectors"][2]) status, job, _ = request("/redact", {"text": REDACT_TEXT, "entities": ["PER", "LOC"]}) assert status == 202 diff --git a/tests/python/test_jobs.py b/tests/python/test_jobs.py index 703526b..5b101f7 100644 --- a/tests/python/test_jobs.py +++ b/tests/python/test_jobs.py @@ -19,7 +19,7 @@ from ai_bridge import jobs, log, tasks from ai_bridge import server as service -from ai_bridge.jobs import CapacityError, JobStore, Settings, TERMINAL +from ai_bridge.jobs import TERMINAL, CapacityError, JobStore, Settings from ai_bridge.server import BridgeServer from ai_bridge.tasks import EMBED_DIMENSIONS, MAX_TEXTS, InvalidInput, _entities, _merge, execute, validate from test_log import parse as parse_log @@ -688,10 +688,10 @@ def test_the_table_covers_every_backend_and_task(self): from ai_bridge.tasks import MODEL_NAMES self.assertEqual(set(MODEL_NAMES), {"lexical", "onnx"}) - for backend, tasks in MODEL_NAMES.items(): + for backend, names in MODEL_NAMES.items(): with self.subTest(backend=backend): - self.assertEqual(set(tasks), {"rerank", "embed", "redact"}) - self.assertTrue(all(isinstance(name, str) and name for name in tasks.values())) + self.assertEqual(set(names), {"rerank", "embed", "redact"}) + self.assertTrue(all(isinstance(name, str) and name for name in names.values())) def test_lexical_results_report_the_names_in_the_table(self): from ai_bridge.tasks import MODEL_NAMES @@ -718,7 +718,7 @@ def test_reports_are_bounded_increasing_and_end_on_total(self): for total in [1, 3, 64, 65, 127, 128, 512]: events = [] execute("rerank", {"query": "x", "documents": ["x"] * total}, "lexical", "", - lambda completed, _: events.append(completed)) + lambda completed, _: events.append(completed)) # noqa: B023 (called only within this iteration) with self.subTest(total=total): self.assertLessEqual(len(events), 66) self.assertEqual(events[0], 0) @@ -781,7 +781,10 @@ def test_hashing_embed_is_explicit_unit_length_and_order_free(self): self.assertAlmostEqual(math.sqrt(sum(x * x for x in vector)), 1.0, places=4) self.assertEqual(vectors[0], vectors[3]) self.assertEqual(vectors[0], vectors[1]) - dot = lambda a, b: sum(x * y for x, y in zip(a, b)) + + def dot(a, b): + return sum(x * y for x, y in zip(a, b)) + self.assertGreater(dot(vectors[0], vectors[1]), dot(vectors[0], vectors[2])) self.assertEqual(done["progress"], {"completed": 4, "total": 4}) @@ -1417,7 +1420,7 @@ def test_listed_tokens_are_absent_from_the_child_environment(self): def serve(bridge, **_): self.store = bridge.jobs - reports.append(self.inspect(value)) + reports.append(self.inspect(value)) # noqa: B023 (called only within this iteration) with patch.dict(os.environ, {"BRIDGE_TEST_TASKS": "1", "BRIDGE_CONCURRENCY": "1", **environment}), \ patch.object(sys, "argv", ["ai_bridge.server", "--port", "0"]), \ diff --git a/tests/python/test_server.py b/tests/python/test_server.py index bc38bea..2c2bd4a 100644 --- a/tests/python/test_server.py +++ b/tests/python/test_server.py @@ -20,7 +20,8 @@ from ai_bridge import jobs, log, metrics, server, tasks from ai_bridge.jobs import Settings from ai_bridge.server import BridgeServer -from test_log import BlockedStream, BrokenStream, parse as parse_log +from test_log import BlockedStream, BrokenStream +from test_log import parse as parse_log TOKEN = "test-only-bridge-token-never-use-in-production" diff --git a/tests/python/test_workflows.py b/tests/python/test_workflows.py index 9633e34..874e83e 100644 --- a/tests/python/test_workflows.py +++ b/tests/python/test_workflows.py @@ -95,7 +95,7 @@ def inputs(step: list) -> list: for line in lines: key = re.match(r"^(\s*)with:\s*$", line) if key: - deeper = itertools.takewhile(lambda below: len(below) - len(below.lstrip()) > len(key.group(1)), lines) + deeper = itertools.takewhile(lambda below: len(below) - len(below.lstrip()) > len(key.group(1)), lines) # noqa: B023 (called only within this iteration) return [below.strip() for below in deeper] return [] diff --git a/worker/ai_bridge/jobs.py b/worker/ai_bridge/jobs.py index c6e084a..0424512 100644 --- a/worker/ai_bridge/jobs.py +++ b/worker/ai_bridge/jobs.py @@ -8,9 +8,9 @@ """ import copy +import json import math import multiprocessing -import json import os import re import signal @@ -19,7 +19,10 @@ import time import uuid from dataclasses import dataclass +from multiprocessing.connection import Connection +from multiprocessing.process import BaseProcess from pathlib import Path +from typing import Any from . import log, metrics, tasks from .tasks import MODEL_NAMES, TEST_TASKS, execute, validate @@ -44,10 +47,12 @@ _INDEX = """DROP INDEX IF EXISTS jobs_idempotency_key; CREATE UNIQUE INDEX IF NOT EXISTS jobs_owner_idempotency_key ON jobs (owner, idempotency_key);""" DEFAULT_OWNER = "default" +# A job's public fields and the store's own, which start with an underscore. +Job = dict[str, Any] class CapacityError(Exception): - def __init__(self, message, retry_after=1): + def __init__(self, message: str, retry_after: int = 1) -> None: super().__init__(message) self.retry_after = retry_after # whole seconds until a submission may be accepted, at least 1 @@ -74,7 +79,7 @@ class Settings: owner_capacity: int | None = None -def _serve(connection, backend, model_dir): +def _serve(connection: Connection, backend: str, model_dir: str) -> None: """A child's whole life: run the jobs the parent sends until it hangs up.""" # Ctrl+C in a terminal reaches the whole process group; stopping children is the parent's job. signal.signal(signal.SIGINT, signal.SIG_IGN) @@ -111,11 +116,11 @@ def _serve(connection, backend, model_dir): except (OSError, ValueError): break - def progress(completed, total): - connection.send((job_id, "progress", {"completed": completed, "total": total})) + def progress(completed: int, total: int) -> None: + connection.send((job_id, "progress", {"completed": completed, "total": total})) # noqa: B023 (called only within this iteration) try: - reply = (job_id, "result", execute(task, payload, backend, model_dir, progress)) + reply: tuple[Any, str, object] = (job_id, "result", execute(task, payload, backend, model_dir, progress)) except Exception: # Exceptions can contain credentials, document contents and local paths. reply = (job_id, "error", "task_failed") @@ -139,21 +144,21 @@ def progress(completed, total): class _Worker: """One child of the pool and what the coordinator knows about it.""" - def __init__(self, process, connection): + def __init__(self, process: BaseProcess, connection: Connection) -> None: self.process = process self.connection = connection self.started_at = time.monotonic() self.ready = False # it has said so, and so reads a job the moment it is sent self.loaded = False # it has loaded every model its backend needs - self.job_id = None # the job sent to it, acknowledged or not + self.job_id: str | None = None # the job sent to it, acknowledged or not self.sending = False # a send outside the lock owns the pipe until it returns self.served = 0 # jobs sent to it, towards max_jobs - self.retired_at = None # set once it must never serve again; it is then reaped outside the lock + self.retired_at: float | None = None # set once it must never serve again; it is then reaped outside the lock self.killed = False class JobStore: - def __init__(self, settings=None): + def __init__(self, settings: Settings | None = None) -> None: self.settings = settings or Settings() if (type(self.settings.concurrency) is not int or type(self.settings.capacity) is not int or type(self.settings.max_jobs) is not int @@ -180,16 +185,16 @@ def __init__(self, settings=None): self._rejected = {(owner, "capacity"): 0 for owner in owners} self._durations = {task: metrics.new_histogram() for task in tasks} self._waits = {task: metrics.new_histogram() for task in tasks} - self._jobs = {} - self._keys = {} # (owner, idempotency key) -> job id, for retained jobs only - self._active = {} # job id -> the worker it was sent to - self._pool = [] # children that may take a job, idle or running one - self._retiring = [] # children that never serve again, until they are gone + self._jobs: dict[str, Job] = {} + self._keys: dict[tuple[str, str | None], str] = {} # (owner, idempotency key) -> job id, for retained jobs only + self._active: dict[str, _Worker] = {} # job id -> the worker it was sent to + self._pool: list[_Worker] = [] # children that may take a job, idle or running one + self._retiring: list[_Worker] = [] # children that never serve again, until they are gone self._backoff_until = 0.0 # no child is started before this time self._lock = threading.RLock() self._stopping = threading.Event() self._draining = False - self._db = None + self._db: sqlite3.Connection | None = None self._open_db() # A child that cannot start now is tried again once the backoff ends; a job that finds no child up # meanwhile fails as worker_start_failed, so the service still answers. @@ -197,7 +202,7 @@ def __init__(self, settings=None): self._thread = threading.Thread(target=self._coordinate, daemon=True) self._thread.start() - def _open_db(self): + def _open_db(self) -> None: """One connection for the store's life, used only under the lock. A file left by an earlier run keeps its retained results; what it still holds as unfinished never completed and is failed, not replayed.""" if self.settings.state_dir is None: @@ -241,7 +246,7 @@ def _open_db(self): if key is not None: self._keys[owner, key] = job_id - def _write(self, job, submitted=False): + def _write(self, job: Job, submitted: bool = False) -> None: """Write a job's row through; a no-op once close() has let go of the connection. A submission that cannot be recorded is refused; a later write that fails leaves this process's copy answering, since the coordinator must not stop over a full disk.""" @@ -265,7 +270,8 @@ def _write(self, job, submitted=False): if submitted: raise CapacityError("job table cannot be written") from None - def submit(self, task, payload, timeout_ms=30000, idempotency_key=None, owner=DEFAULT_OWNER): + def submit(self, task: str, payload: Any, timeout_ms: int = 30000, idempotency_key: str | None = None, + owner: str = DEFAULT_OWNER) -> Job: if owner not in self.settings.owners: raise ValueError("owner is not one of the service's token names") if not isinstance(task, str): @@ -284,7 +290,7 @@ def submit(self, task, payload, timeout_ms=30000, idempotency_key=None, owner=DE self._rejected[owner, "capacity"] += 1 raise - def _accept(self, task, payload, timeout_ms, idempotency_key, owner): + def _accept(self, task: str, payload: Any, timeout_ms: int, idempotency_key: str | None, owner: str) -> Job: if self._draining: raise DrainingError("worker is draining") if self._stopping.is_set(): @@ -315,18 +321,18 @@ def _accept(self, task, payload, timeout_ms, idempotency_key, owner): self._submitted[owner, task] += 1 return self._public(job) - def _owned(self, job_id, owner): + def _owned(self, job_id: str, owner: str) -> Job | None: """The job if this owner submitted it; another owner's job is as unknown to it as one that never was.""" job = self._jobs.get(job_id) return job if job is not None and job["_owner"] == owner else None - def get(self, job_id, owner=DEFAULT_OWNER): + def get(self, job_id: str, owner: str = DEFAULT_OWNER) -> Job | None: with self._lock: self._expire() job = self._owned(job_id, owner) return None if job is None else self._public(job) - def cancel(self, job_id, owner=DEFAULT_OWNER): + def cancel(self, job_id: str, owner: str = DEFAULT_OWNER) -> Job | None: with self._lock: self._expire() job = self._owned(job_id, owner) @@ -337,10 +343,10 @@ def cancel(self, job_id, owner=DEFAULT_OWNER): return self._public(job) @property - def draining(self): + def draining(self) -> bool: return self._draining - def drain(self): + def drain(self) -> None: """Take no new job and fail every job no child has started, dropping its input; running jobs go on.""" with self._lock: self._draining = True @@ -351,12 +357,12 @@ def drain(self): if job["status"] == "queued": self._abort(job, "failed", code="service_restarted") - def busy(self): + def busy(self) -> bool: """Whether any job has yet to end.""" with self._lock: return any(job["status"] not in TERMINAL for job in self._jobs.values()) - def readiness(self): + def readiness(self) -> str | None: """None when a child is idle with its models loaded; otherwise why not, from a fixed set.""" with self._lock: if self._draining: @@ -366,7 +372,7 @@ def readiness(self): return None return "models_loading" if idle else "no_idle_child" - def snapshot(self): + def snapshot(self) -> dict[str, Any]: """The aggregates for GET /metrics as plain values, copied under the lock so the caller formats them without holding it.""" with self._lock: @@ -380,19 +386,19 @@ def snapshot(self): "duration": copy.deepcopy(self._durations), "queue": copy.deepcopy(self._waits), } - def _retry_after(self, owner=None): + def _retry_after(self, owner: str | None = None) -> int: """Seconds until the first retained job (of this owner, when given) expires; holding only unfinished jobs gives a short default, since none of them has a known end.""" - finished = [job["_finished"] for job in self._jobs.values() + finished: list[float] = [job["_finished"] for job in self._jobs.values() if job["_finished"] is not None and owner in (None, job["_owner"])] if not finished: return 1 return max(1, math.ceil(min(finished) + self.settings.retention_seconds - time.monotonic())) - def _public(self, job): + def _public(self, job: Job) -> Job: return copy.deepcopy({key: value for key, value in job.items() if not key.startswith("_")}) - def _expire(self): + def _expire(self) -> None: now = time.monotonic() expired = [] for job_id in list(self._jobs): @@ -408,7 +414,7 @@ def _expire(self): except sqlite3.Error: pass # the next open deletes them by their finish time - def _finish(self, job, status, result=None, code=None): + def _finish(self, job: Job, status: str, result: Any = None, code: str | None = None) -> None: job.update(status=status, result=result, error={"code": code, "message": "Task did not complete"} if code else None) job["_input"] = None job["_finished"] = time.monotonic() @@ -422,7 +428,7 @@ def _finish(self, job, status, result=None, code=None): queue_ms=None if started is None else round((started - job["_submitted_at"]) * 1000, 3), run_ms=None if started is None else round((job["_finished"] - started) * 1000, 3)) - def _abort(self, job, status, code=None): + def _abort(self, job: Job, status: str, code: str | None = None) -> None: """End a job before its child has: the child is retired, since it may be halfway through the input.""" worker = self._active.pop(job["id"], None) if worker is not None: @@ -430,7 +436,7 @@ def _abort(self, job, status, code=None): self._retire(worker) self._finish(job, status, code=code) - def _start_worker(self): + def _start_worker(self) -> _Worker | None: """A fresh child, or None when the system refuses one; nothing allocated is leaked either way. Runs without the lock, since creating a process takes milliseconds on Linux and tens of them on Windows.""" parent = child = process = None @@ -453,7 +459,7 @@ def _start_worker(self): child.close() return _Worker(process, parent) - def _start_failed(self): + def _start_failed(self) -> None: """A child could not be started, or died before it was ready: wait before trying again, and when no child is up to take them, tell the queued jobs now rather than at their deadline.""" self._backoff_until = time.monotonic() + BACKOFF_SECONDS @@ -463,7 +469,7 @@ def _start_failed(self): if job["status"] == "queued" and job["id"] not in self._active: self._finish(job, "failed", code="worker_start_failed") - def _replenish(self): + def _replenish(self) -> None: """Fill the pool up to the concurrency limit, deciding and adding under the lock and starting without it. Only the coordinator calls it once the store is built, so no two starts race for the same place.""" while True: @@ -484,7 +490,7 @@ def _replenish(self): return @staticmethod - def _discard(worker): + def _discard(worker: _Worker) -> None: """Stop a child that never joined the pool.""" try: worker.connection.close() @@ -500,7 +506,7 @@ def _discard(worker): except ValueError: pass - def _retire(self, worker, stop=False, died=False): + def _retire(self, worker: _Worker, stop: bool = False, died: bool = False) -> None: """Take a child out of service for good. With stop, it is asked to exit by itself, because it is idle and has done its jobs; otherwise it is terminated, since it may be running one. Either way the reaper kills it if it is still there after the grace period, outside the lock. A child that died (or went silent) before @@ -510,7 +516,8 @@ def _retire(self, worker, stop=False, died=False): if worker.retired_at is not None: return job = self._jobs.get(worker.job_id) if worker.job_id is not None else None - self._active.pop(worker.job_id, None) + if worker.job_id is not None: + self._active.pop(worker.job_id, None) worker.job_id = None # A send in flight owns the pipe: the child is terminated instead, and the sender closes the pipe. stop = stop and not worker.sending @@ -539,14 +546,14 @@ def _retire(self, worker, stop=False, died=False): else: self._finish(job, "failed", code="worker_crashed") - def _assign(self, job): + def _assign(self, job: Job) -> None: """Send a queued job to an idle child that is ready for one; otherwise it waits for one, which the coordinator starts outside the lock when the pool has room.""" attached = self._attach(job) if attached is not None: self._send(*attached) - def _attach(self, job): + def _attach(self, job: Job) -> tuple[_Worker, tuple[str, str, Any]] | None: """Under the lock: give a queued job to an idle ready child; returns (worker, message) to send, or None.""" worker = next((candidate for candidate in self._pool if candidate.ready and candidate.job_id is None), None) if worker is None: @@ -557,7 +564,7 @@ def _attach(self, job): self._active[job["id"]] = worker return worker, (job["id"], job["task"], job["_input"]) - def _send(self, worker, message): + def _send(self, worker: _Worker, message: tuple[str, str, Any]) -> None: """Send what _attach returned; the coordinator calls it without the lock, so a full pipe never holds a GET.""" try: worker.connection.send(message) @@ -575,9 +582,9 @@ def _send(self, worker, message): # Its pipe is gone, so the child is: retired like one that died, and replaced. self._retire(worker) - def _coordinate(self): + def _coordinate(self) -> None: while not self._stopping.wait(0.01): - sends = [] + sends: list[tuple[_Worker, tuple[str, str, Any]]] = [] with self._lock: self._expire() now = time.monotonic() @@ -611,7 +618,7 @@ def _coordinate(self): return # close() reaps from here on; reaping too would touch processes it has closed self._reap(retiring) - def _drain(self, worker): + def _drain(self, worker: _Worker) -> None: """Deliver what the child has sent about the job it runs, up to that job's end.""" try: while worker.connection.poll(): @@ -621,7 +628,7 @@ def _drain(self, worker): # The child is gone (end of file), or sent something that is not its protocol: either way it is done. self._retire(worker, died=True) - def _drain_message(self, worker, message): + def _drain_message(self, worker: _Worker, message: tuple[Any, str, Any]) -> bool: """Apply one message from a child; True once it ended the job the child was running.""" job_id, kind, value = message if kind == "ready": @@ -653,11 +660,12 @@ def _drain_message(self, worker, message): self._retire(worker, stop=True) return True - def _reap(self, retiring): + def _reap(self, retiring: list[_Worker]) -> None: """Outside the lock: let go of retired children that have exited, and kill those that have not in time.""" - gone = [] + gone: list[_Worker] = [] now = time.monotonic() for worker in retiring: + assert worker.retired_at is not None # set by _retire before the worker joins the list if worker.process.is_alive(): if not worker.killed and now - worker.retired_at >= GRACE_SECONDS: worker.process.kill() @@ -672,7 +680,7 @@ def _reap(self, retiring): if worker in self._retiring: # close() may have cleared the list meanwhile self._retiring.remove(worker) - def close(self): + def close(self) -> int: """Stop every child; returns how many unfinished jobs this ended as cancelled.""" self._stopping.set() self._thread.join(timeout=2) diff --git a/worker/ai_bridge/log.py b/worker/ai_bridge/log.py index 102d380..2cc0cc0 100644 --- a/worker/ai_bridge/log.py +++ b/worker/ai_bridge/log.py @@ -12,6 +12,8 @@ import sys import threading import time +from collections.abc import Callable +from typing import Any, TextIO, cast from .tasks import MODEL_NAMES, TEST_TASKS @@ -35,24 +37,24 @@ QUEUE_LINES = 1024 -def _count(value): +def _count(value: object) -> bool: return type(value) is int and value >= 0 -def _duration(value): - return type(value) in (int, float) and math.isfinite(value) and value >= 0 +def _duration(value: object) -> bool: + return type(value) in (int, float) and math.isfinite(cast(float, value)) and cast(float, value) >= 0 -def _matches(pattern): +def _matches(pattern: re.Pattern[str]) -> Callable[[object], bool]: return lambda value: type(value) is str and pattern.fullmatch(value) is not None -def _among(choices): +def _among(choices: Any) -> Callable[[object], bool]: return lambda value: type(value) is str and value in choices # What each key's value must be; anything else is written as null. -_VALID = { +_VALID: dict[str, Callable[[object], bool]] = { "method": _among(METHODS), "route": _among(ROUTES), "status": lambda value: type(value) is int and 100 <= value <= 599, "duration_ms": _duration, "queue_ms": _duration, "run_ms": _duration, @@ -64,11 +66,11 @@ def _among(choices): } -def _event(record): +def _event(record: logging.LogRecord) -> tuple[str, dict[str, Any]] | None: """The (event, fields) a record carries, or None for a record event() did not make.""" carried = getattr(record, "bridge", None) if type(carried) is tuple and len(carried) == 2 and carried[0] in EVENTS and type(carried[1]) is dict: - return carried + return cast(tuple[str, dict[str, Any]], carried) return None @@ -76,10 +78,13 @@ class JsonFormatter(logging.Formatter): """One JSON object from the record's event and fields alone: the message, its arguments, an exception and a stack are never read.""" - def format(self, record): - name, fields = _event(record) + def format(self, record: logging.LogRecord) -> str: + carried = _event(record) + # The handler's filter passes only the records event() made. + assert carried is not None + name, fields = carried moment = time.gmtime(record.created) - line = { + line: dict[str, object] = { "ts": time.strftime("%Y-%m-%dT%H:%M:%S", moment) + ".%03dZ" % int(record.msecs), "level": "warning" if record.levelno >= logging.WARNING else "info", "event": name, @@ -93,24 +98,24 @@ def format(self, record): class _Enqueue(logging.Handler): """Formats in the caller and hands the line over without waiting; a full queue drops it.""" - def __init__(self, lines): + def __init__(self, lines: queue.Queue[str | None]) -> None: super().__init__() self._lines = lines self.setFormatter(JsonFormatter()) self.addFilter(lambda record: _event(record) is not None) - def emit(self, record): + def emit(self, record: logging.LogRecord) -> None: try: self._lines.put_nowait(self.format(record)) except Exception: pass - def handleError(self, record): + def handleError(self, record: logging.LogRecord) -> None: # The standard handler prints the exception to stderr, which is what must never happen here. pass -def _write(lines, stream, stopped): +def _write(lines: queue.Queue[str | None], stream: TextIO, stopped: threading.Event) -> None: while True: try: line = lines.get(timeout=0.1) @@ -134,10 +139,10 @@ def _write(lines, stream, stopped): # Without a handler of its own the standard library's last resort would print warnings to stderr. _logger.addHandler(logging.NullHandler()) _lock = threading.Lock() -_active = None # (handler, queue, writer thread, stop event) while json is on +_active: tuple[_Enqueue, queue.Queue[str | None], threading.Thread, threading.Event] | None = None # while json is on -def configure(mode, stream=None): +def configure(mode: str, stream: TextIO | None = None) -> None: """Start writing JSON lines to stream (stderr when None) for json; off writes nothing. Replaces an earlier one.""" global _active if mode not in MODES: @@ -145,7 +150,7 @@ def configure(mode, stream=None): close() if mode == "off": return - lines = queue.Queue(QUEUE_LINES) + lines: queue.Queue[str | None] = queue.Queue(QUEUE_LINES) stopped = threading.Event() writer = threading.Thread(target=_write, args=(lines, sys.stderr if stream is None else stream, stopped), daemon=True) @@ -156,7 +161,7 @@ def configure(mode, stream=None): _logger.addHandler(handler) -def close(): +def close() -> None: """Stop logging and let the writer finish what is queued, waiting a moment at most for a stream that blocks.""" global _active with _lock: @@ -173,7 +178,7 @@ def close(): writer.join(1) -def event(name, warning=False, **fields): +def event(name: str, warning: bool = False, **fields: object) -> None: """Write one line of this event; a key outside its schema is dropped and a value of the wrong kind is null.""" if _active is None: return diff --git a/worker/ai_bridge/metrics.py b/worker/ai_bridge/metrics.py index 69f6e1c..3d3805e 100644 --- a/worker/ai_bridge/metrics.py +++ b/worker/ai_bridge/metrics.py @@ -4,58 +4,62 @@ import bisect import math +from collections.abc import Iterable +from typing import Any + +Labels = dict[str, str] CONTENT_TYPE = "text/plain; version=0.0.4; charset=utf-8" # Upper bounds in seconds; a job longer than the last falls only into +Inf. BUCKETS = (0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10, 30, 60, 300) -def new_histogram(): +def new_histogram() -> dict[str, Any]: """Counts per bucket, not cumulative, the last for +Inf only; render() adds them up.""" return {"buckets": [0] * (len(BUCKETS) + 1), "sum": 0.0, "count": 0} -def observe(histogram, seconds): +def observe(histogram: dict[str, Any], seconds: float) -> None: # bisect_left puts a value equal to a bound in that bound's bucket, as le (less or equal) means. histogram["buckets"][bisect.bisect_left(BUCKETS, seconds)] += 1 histogram["sum"] += seconds histogram["count"] += 1 -def escape(value): +def escape(value: str) -> str: return value.replace("\\", "\\\\").replace('"', '\\"').replace("\n", "\\n") -def number(value): +def number(value: float) -> str: if isinstance(value, float): return "+Inf" if value == math.inf else repr(value) return str(value) -def _labels(labels): +def _labels(labels: Labels) -> str: if not labels: return "" return "{" + ",".join(f'{key}="{escape(value)}"' for key, value in labels.items()) + "}" -def _family(name, kind, help_text, lines): +def _family(name: str, kind: str, help_text: str, lines: Iterable[str]) -> str: return f"# HELP {name} {help_text}\n# TYPE {name} {kind}\n" + "".join(line + "\n" for line in lines) -def _simple(name, kind, help_text, samples): +def _simple(name: str, kind: str, help_text: str, samples: Iterable[tuple[Labels, float]]) -> str: return _family(name, kind, help_text, [f"{name}{_labels(labels)} {number(value)}" for labels, value in samples]) -def counter(name, help_text, samples): +def counter(name: str, help_text: str, samples: Iterable[tuple[Labels, float]]) -> str: return _simple(name, "counter", help_text, samples) -def gauge(name, help_text, samples): +def gauge(name: str, help_text: str, samples: Iterable[tuple[Labels, float]]) -> str: return _simple(name, "gauge", help_text, samples) -def histogram(name, help_text, series): - lines = [] +def histogram(name: str, help_text: str, series: Iterable[tuple[Labels, dict[str, Any]]]) -> str: + lines: list[str] = [] for labels, values in series: total = 0 for bound, count in zip(BUCKETS + (math.inf,), values["buckets"]): @@ -66,7 +70,7 @@ def histogram(name, help_text, series): return _family(name, "histogram", help_text, lines) -def render(snapshot): +def render(snapshot: dict[str, Any]) -> str: """The whole body from JobStore.snapshot(), which is a copy: nothing here holds the store's lock.""" tasks = snapshot["tasks"] return "".join([ diff --git a/worker/ai_bridge/server.py b/worker/ai_bridge/server.py index d0c1c10..43eb910 100644 --- a/worker/ai_bridge/server.py +++ b/worker/ai_bridge/server.py @@ -2,6 +2,7 @@ import argparse import hmac +import io import json import os import re @@ -11,8 +12,10 @@ import stat import threading import time +from collections.abc import Callable, Iterable from dataclasses import replace from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Any from . import log from .jobs import DEFAULT_OWNER, CapacityError, DrainingError, JobStore, Settings @@ -39,12 +42,12 @@ class _Head: """The connection's reader, which refuses a request line or header block over its limit before reading it all; the body is read through it unchanged.""" - def __init__(self, rfile): + def __init__(self, rfile: io.BufferedIOBase) -> None: self._rfile = rfile - self._left = None # bytes the header block may still take; None until the request line is read + self._left: int | None = None # bytes the header block may still take; None until the request line is read self._done = False - def readline(self, limit=-1): + def readline(self, limit: int = -1) -> bytes: if self._done: return self._rfile.readline(limit) if self._left is None: @@ -65,14 +68,14 @@ def readline(self, limit=-1): raise HeadTooLarge return line - def __getattr__(self, name): + def __getattr__(self, name: str) -> Any: return getattr(self._rfile, name) MAX_TOKENS = 64 MAX_TOKENS_FILE = 65536 TOKEN_NAME = re.compile(r"[A-Za-z0-9_.-]{1,64}") -def check_tokens(tokens, source="the token list"): +def check_tokens(tokens: dict[str, str], source: str = "the token list") -> dict[str, str]: """Refuse a token list the service must not start with; the message names the source, never a name or token.""" if not tokens: raise ValueError(f"{source} holds no token") @@ -88,11 +91,11 @@ def check_tokens(tokens, source="the token list"): return tokens -def _entries(lines, source): +def _entries(lines: list[str], source: str) -> dict[str, str]: """name=token entries as a mapping; the token is everything after the first "=", since a token may hold one.""" if len(lines) > MAX_TOKENS: raise ValueError(f"{source} holds more than {MAX_TOKENS} tokens") - tokens = {} + tokens: dict[str, str] = {} for line in lines: name, separator, token = line.partition("=") if not separator or not TOKEN_NAME.fullmatch(name): @@ -104,7 +107,7 @@ def _entries(lines, source): return check_tokens(tokens, source) -def _read_tokens_file(path): +def _read_tokens_file(path: str) -> dict[str, str]: """The file's entries, read once; a file others could read or replace is refused, since it holds every token.""" try: # Non-blocking, so a FIFO at the path cannot hold the start up; a regular file reads as usual. @@ -128,7 +131,7 @@ def _read_tokens_file(path): return _entries(lines, "BRIDGE_TOKENS_FILE") -def resolve_tokens(): +def resolve_tokens() -> dict[str, str]: """The service's tokens by name from BRIDGE_TOKEN (named default), BRIDGE_TOKENS or BRIDGE_TOKENS_FILE, exactly one of them. Both token variables leave the environment here, before any child exists: the store's children inherit the environment block, and a variable removed later stays in /proc//environ. The file's path may stay.""" @@ -155,7 +158,8 @@ class BridgeServer(ThreadingHTTPServer): daemon_threads = True request_queue_size = 16 - def __init__(self, address, tokens, settings=None, drain_seconds=30): + def __init__(self, address: tuple[str, int], tokens: str | dict[str, str], settings: Settings | None = None, + drain_seconds: int = 30) -> None: """tokens maps each name to its token; a single token is the name default.""" if isinstance(tokens, str): tokens = {DEFAULT_OWNER: tokens} @@ -164,9 +168,9 @@ def __init__(self, address, tokens, settings=None, drain_seconds=30): self.credentials = tuple((name, ("Bearer " + token).encode("ascii")) for name, token in tokens.items()) self.drain_seconds = drain_seconds self.cancelled = 0 # unfinished jobs the stop ended, for the stop line - self._signalled_at = None + self._signalled_at: float | None = None self._interrupted = threading.Event() - self._drainer = None + self._drainer: threading.Thread | None = None self.slots = threading.BoundedSemaphore(32) # The children the store starts inherit this process's environment block as it is at that moment, and # removing a variable later leaves it in /proc//environ; so the tokens leave before any child exists. @@ -179,7 +183,7 @@ def __init__(self, address, tokens, settings=None, drain_seconds=30): self.jobs.close() raise - def process_request(self, request, address): + def process_request(self, request: Any, address: Any) -> None: if not self.slots.acquire(blocking=False): # One non-blocking send that fits any socket buffer, nothing read: a client cannot hold the accept loop. try: @@ -196,24 +200,24 @@ def process_request(self, request, address): self.slots.release() raise - def process_request_thread(self, request, address): + def process_request_thread(self, request: Any, address: Any) -> None: try: super().process_request_thread(request, address) finally: self.slots.release() - def handle_error(self, request, client_address): + def handle_error(self, request: Any, client_address: Any) -> None: # The default prints the peer's address and a traceback to stderr, which takes only JSON lines or nothing. pass - def server_close(self): + def server_close(self) -> None: super().server_close() if self._drainer is not None: # A second signal stops serving before the drain thread has written its line, which goes before stop's. self._drainer.join(1) self.cancelled = self.jobs.close() - def stop(self): + def stop(self) -> None: """A signal's handler: the first starts the drain, a second ends it at once.""" if self._signalled_at is not None: self._interrupted.set() @@ -224,8 +228,9 @@ def stop(self): self._drainer = threading.Thread(target=self._drain_then_stop, daemon=True) self._drainer.start() - def _drain_then_stop(self): + def _drain_then_stop(self) -> None: """Until running jobs end or the deadline passes, then stop serving; server_close ends what still runs.""" + assert self._signalled_at is not None # stop() sets it before it starts this thread self.jobs.drain() deadline = self._signalled_at + self.drain_seconds while not self._interrupted.is_set() and self.jobs.busy(): @@ -241,15 +246,20 @@ def _drain_then_stop(self): class Handler(BaseHTTPRequestHandler): server_version = "PhpPythonBridge/0.1" sys_version = "" + server: BridgeServer + _owner: str | None + _status: int | None - def setup(self): + def setup(self) -> None: super().setup() self.connection.settimeout(5) - self.rfile = _Head(self.rfile) + # _Head is not a BufferedIOBase: it forwards every other attribute to the stream it wraps. + self.rfile = _Head(self.rfile) # type: ignore[assignment] - def handle_one_request(self): + def handle_one_request(self) -> None: started = time.monotonic() - self.command = self.path = self._owner = self._status = None + # The standard library types command and path as str; they are None until a request line parses. + self.command = self.path = self._owner = self._status = None # type: ignore[assignment] try: super().handle_one_request() except HeadTooLarge: @@ -263,12 +273,12 @@ def handle_one_request(self): status=self._status, duration_ms=round((time.monotonic() - started) * 1000, 3), owner=self._owner) - def _method(self): + def _method(self) -> str | None: if self.command is None: return None return self.command if self.command in ("GET", "POST") else "other" - def _route(self): + def _route(self) -> str: if self.path in _ROUTES: return _ROUTES[self.path] if self.path is not None: @@ -277,23 +287,23 @@ def _route(self): return route return "other" - def send_response(self, code, message=None): + def send_response(self, code: int, message: str | None = None) -> None: # The standard library's own refusals pass an HTTPStatus. self._status = int(code) super().send_response(code, message) - def log_message(self, *_): + def log_message(self, *_: Any) -> None: # Do not log URLs, credentials, inputs or model results. pass - def _reply(self, status, value, headers=()): + def _reply(self, status: int, value: object, headers: Iterable[tuple[str, str]] = ()) -> None: self._send(status, json.dumps(value, allow_nan=False, separators=(",", ":")).encode(), "application/json", headers) - def _text(self, status, text): + def _text(self, status: int, text: str) -> None: self._send(status, text.encode("utf-8"), CONTENT_TYPE) - def _send(self, status, body, content_type, headers=()): + def _send(self, status: int, body: bytes, content_type: str, headers: Iterable[tuple[str, str]] = ()) -> None: self.send_response(status) for name, header in headers: self.send_header(name, header) @@ -305,10 +315,10 @@ def _send(self, status, body, content_type, headers=()): self.close_connection = True self.wfile.write(body) - def _error(self, status, code, message, headers=()): + def _error(self, status: int, code: str, message: str, headers: Iterable[tuple[str, str]] = ()) -> None: self._reply(status, {"error": {"code": code, "message": message}}, headers) - def _authorized(self): + def _authorized(self) -> str | None: """The name of the token presented, or None once a 401 is sent. Every token is compared, whichever matches.""" values = self.headers.get_all("Authorization", []) presented = values[0].encode("utf-8") if len(values) == 1 else b"" @@ -322,7 +332,7 @@ def _authorized(self): self._owner = owner return owner - def _body(self): + def _body(self) -> dict[str, Any]: lengths = self.headers.get_all("Content-Length", []) if self.headers.get("Transfer-Encoding") or len(lengths) != 1 or not re.fullmatch(r"[0-9]{1,7}", lengths[0]): raise ValueError("A single Content-Length is required; chunked bodies are unsupported") @@ -335,11 +345,11 @@ def _body(self): if len(raw) != size: raise ValueError("Incomplete request body") - def reject_constant(_): + def reject_constant(_: str) -> None: raise ValueError("Non-finite JSON numbers are unsupported") - def unique_keys(pairs): - result = {} + def unique_keys(pairs: list[tuple[str, Any]]) -> dict[str, Any]: + result: dict[str, Any] = {} for key, value in pairs: if key in result: raise ValueError("Duplicate JSON keys are unsupported") @@ -351,7 +361,7 @@ def unique_keys(pairs): raise ValueError("JSON body must be an object") return body - def do_GET(self): + def do_GET(self) -> None: owner = self._authorized() if owner is None: return @@ -376,7 +386,7 @@ def do_GET(self): else: self._reply(200, job) - def do_POST(self): + def do_POST(self) -> None: owner = self._authorized() if owner is None: return @@ -391,7 +401,7 @@ def do_POST(self): # The store reads None as no key, so an explicit null is refused here rather than ignored. if "idempotency_key" in body and body["idempotency_key"] is None: raise ValueError("idempotency_key must be a string") - job = self.server.jobs.submit(body["task"], body["input"], body.get("timeout_ms", 30000), + job: dict[str, Any] | None = self.server.jobs.submit(body["task"], body["input"], body.get("timeout_ms", 30000), body.get("idempotency_key"), owner) self._reply(202, job) return @@ -416,7 +426,7 @@ def do_POST(self): self._error(408, "request_timeout", "Request body was not received in time") -def _number(name, default, kind): +def _number[T](name: str, default: str, kind: Callable[[str], T]) -> T: """A numeric setting from the environment; a malformed one stops the service with the variable's name.""" try: return kind(os.environ.get(name, default)) @@ -424,7 +434,7 @@ def _number(name, default, kind): raise SystemExit(f"{name} must be {'an integer' if kind is int else 'a number'}") from None -def main(): +def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--host", default="127.0.0.1") parser.add_argument("--port", type=int, default=8090) @@ -439,7 +449,7 @@ def main(): test_tasks=os.environ.get("BRIDGE_TEST_TASKS") == "1", max_jobs=_number("BRIDGE_WORKER_MAX_JOBS", "1000", int), state_dir=os.environ.get("BRIDGE_STATE_DIR") or None, - owner_capacity=_number("BRIDGE_OWNER_CAPACITY", None, int) if "BRIDGE_OWNER_CAPACITY" in os.environ else None, + owner_capacity=_number("BRIDGE_OWNER_CAPACITY", "", int) if "BRIDGE_OWNER_CAPACITY" in os.environ else None, ) if settings.owner_capacity is not None and not 1 <= settings.owner_capacity <= settings.capacity: raise SystemExit(f"BRIDGE_OWNER_CAPACITY must be an integer from 1 to BRIDGE_CAPACITY ({settings.capacity})") @@ -467,7 +477,7 @@ def main(): retention_seconds=settings.retention_seconds, max_jobs=settings.max_jobs, drain_seconds=drain_seconds, tokens=len(tokens)) with bridge as server: - def stop(*_): + def stop(*_: Any) -> None: server.stop() signal.signal(signal.SIGTERM, stop) diff --git a/worker/ai_bridge/tasks.py b/worker/ai_bridge/tasks.py index 2d54613..9af5f6b 100644 --- a/worker/ai_bridge/tasks.py +++ b/worker/ai_bridge/tasks.py @@ -6,8 +6,10 @@ import os import re import time +from collections.abc import Callable, Iterator from pathlib import Path - +from types import FrameType +from typing import Any MAX_DOCUMENTS = 512 # Query-document pairs per forward pass. An estimate, not a measurement: at 512 tokens TinyBERT-L2 keeps @@ -42,6 +44,10 @@ # Present only with BRIDGE_TEST_TASKS=1, for the fault tests. TEST_TASKS = ("test.delay", "test.crash", "test.error", "test.inspect") +# Called with (completed, total) as a task advances. +Progress = Callable[[int, int], None] +# A found span of text, as (start, end) offsets. +Finder = Callable[[str], Iterator[tuple[int, int]]] class InvalidInput(ValueError): @@ -66,7 +72,7 @@ def _strings(values: object, plural: str, singular: str, limit: int, extra: int raise InvalidInput(f"{plural} must total at most {MAX_TOTAL_CHARACTERS} characters") -def validate(task: str, payload: object, test_tasks: bool = False) -> dict: +def validate(task: str, payload: object, test_tasks: bool = False) -> dict[str, Any]: if not isinstance(payload, dict): raise InvalidInput("input must be an object") if test_tasks and task in TEST_TASKS: @@ -111,7 +117,7 @@ def validate(task: str, payload: object, test_tasks: bool = False) -> dict: return payload -def execute(task: str, payload: dict, backend: str, model_dir: str, progress) -> dict: +def execute(task: str, payload: dict[str, Any], backend: str, model_dir: str, progress: Progress) -> dict[str, Any]: if task == "test.crash": os._exit(17) if task == "test.error": @@ -134,7 +140,7 @@ def execute(task: str, payload: dict, backend: str, model_dir: str, progress) -> # cannot flood the coordinator pipe. step = -(-total // 64) - def report(completed): + def report(completed: int) -> None: if completed % step == 0 or completed == total: progress(completed, total) @@ -161,7 +167,7 @@ def report(completed): return {"rankings": rankings[:payload.get("top_k", total)], "model": model} -def _reachable_holding(needle: str, roots: list) -> int: +def _reachable_holding(needle: str, roots: list[object]) -> int: """How many distinct strings reachable from the roots contain the needle, following every reference the collector knows of, the needle object itself excepted. The collector does not track a dict or tuple of plain strings, so such a container is reached through whatever holds it rather than listed on its own; the walk @@ -183,7 +189,7 @@ def _reachable_holding(needle: str, roots: list) -> int: return held -def _inspect(parts: list) -> dict: +def _inspect(parts: list[str]) -> dict[str, Any]: """What a test may ask a worker child about itself: its process id, whether a service token variable or the text joined from parts reached its environment or its initial environment block (which unsetenv leaves in place), and how many strings reachable from any live object or any frame's locals still hold that text, which no earlier job @@ -199,7 +205,9 @@ def _inspect(parts: list) -> dict: # cached on the frame and brought up to date only when it is read, so a read here drops what an earlier # inspection left in it before the collector's objects are listed. The own frame is known by id, not # held: a frame object still referenced when its function returns takes over the locals. - roots, own = [], id(sys._getframe()) + roots: list[object] = [] + own = id(sys._getframe()) + frame: FrameType | None for frame in sys._current_frames().values(): while frame is not None: if id(frame) != own: @@ -208,6 +216,8 @@ def _inspect(parts: list) -> dict: gc.collect() roots.extend(gc.get_objects()) held = _reachable_holding(needle, roots) + in_block: bool | None + needle_in_block: bool | None try: with open("/proc/self/environ", "rb") as environ: block = environ.read().split(b"\0") @@ -221,7 +231,7 @@ def _inspect(parts: list) -> dict: "needle_in_environment_block": needle_in_block} -def _tensors(encodings: list) -> dict: +def _tensors(encodings: list[Any]) -> dict[str, Any]: """The model inputs for a list of encodings of equal length, as int64 arrays. The only numpy touchpoint.""" import numpy as np @@ -232,8 +242,8 @@ def _tensors(encodings: list) -> dict: } -def _score_pairs(tokenizer, session, names: set, query: str, documents: list, report, - batch_size: int = RERANK_BATCH_SIZE) -> list: +def _score_pairs(tokenizer: Any, session: Any, names: set[str], query: str, documents: list[str], + report: Callable[[int], None], batch_size: int = RERANK_BATCH_SIZE) -> list[float]: """One cross-encoder forward pass per batch of (query, document) pairs, scores in document order. The tokenizer pads each batch to its own longest pair and the attention mask hides the padding, so a @@ -241,7 +251,7 @@ def _score_pairs(tokenizer, session, names: set, query: str, documents: list, re """ if batch_size < 1: raise ValueError("batch size must be positive") - scores = [] + scores: list[float] = [] for start in range(0, len(documents), batch_size): batch = documents[start:start + batch_size] inputs = _tensors(tokenizer.encode_batch([(query, document) for document in batch])) @@ -257,17 +267,17 @@ def _score_pairs(tokenizer, session, names: set, query: str, documents: list, re # Loaded tokenizers and sessions by model directory and tokenizer settings. A worker child loads each model it needs # once, on first use, and keeps it for its next jobs; the coordinator never runs a task, so it holds none. -_LOADED = {} +_LOADED: dict[tuple[str, int, int], tuple[Any, Any, set[str]]] = {} -def _onnx(directory: Path, max_length: int, stride: int = 0): +def _onnx(directory: Path, max_length: int, stride: int = 0) -> tuple[Any, Any, set[str]]: key = (str(directory), max_length, stride) if key not in _LOADED: _LOADED[key] = _load_onnx(directory, max_length, stride) return _LOADED[key] -def _load_onnx(directory: Path, max_length: int, stride: int): +def _load_onnx(directory: Path, max_length: int, stride: int) -> tuple[Any, Any, set[str]]: # Optional dependencies are installed separately; no runtime downloads. import onnxruntime as ort from tokenizers import Tokenizer @@ -284,13 +294,13 @@ def _load_onnx(directory: Path, max_length: int, stride: int): return tokenizer, session, {item.name for item in session.get_inputs()} -def missing_models(model_dir: str) -> list: +def missing_models(model_dir: str) -> list[str]: """The model directories under model_dir that lack either file, in MODEL_DIRECTORIES order.""" return [name for name in MODEL_DIRECTORIES if not all((Path(model_dir) / name / file).is_file() for file in ("tokenizer.json", "model.onnx"))] -def _signed_counts(words: list) -> list: +def _signed_counts(words: list[str]) -> list[float]: # sha256 keeps the demo stable across processes; hash() is salted per process. vector = [0.0] * EMBED_DIMENSIONS for word in words: @@ -299,7 +309,7 @@ def _signed_counts(words: list) -> list: return vector -def _hashing_vector(text: str) -> list: +def _hashing_vector(text: str) -> list[float]: vector = _signed_counts(re.findall(r"\w+", text.casefold()) or [text.strip()]) if not any(vector): # Opposite signs in a shared bucket can cancel exactly; a single token never does. @@ -308,7 +318,7 @@ def _hashing_vector(text: str) -> list: return [value / norm for value in vector] -def _embed(texts: list, backend: str, model_dir: str, progress) -> dict: +def _embed(texts: list[str], backend: str, model_dir: str, progress: Progress) -> dict[str, Any]: total = len(texts) progress(0, total) if backend == "onnx": @@ -337,7 +347,7 @@ def _embed(texts: list, backend: str, model_dir: str, progress) -> dict: return {"model": model, "dimensions": EMBED_DIMENSIONS, "vectors": vectors} -def _output(session, outputs: list, name: str): +def _output(session: Any, outputs: list[Any], name: str) -> Any: # Take an output by name; position is only a fallback for exports without names. names = [item.name for item in session.get_outputs()] return outputs[names.index(name)] if name in names else outputs[0] @@ -353,7 +363,7 @@ def _luhn(digits: str) -> bool: return total % 10 == 0 -def _card_range(text: str, groups: list, start: int): +def _card_range(text: str, groups: list[tuple[int, int]], start: int) -> tuple[int, int] | None: # Longest range first, then leftmost, so the most conservative valid mask wins. ranges = [] for first in range(start, len(groups)): @@ -369,7 +379,7 @@ def _card_range(text: str, groups: list, start: int): return None -def _cards(text: str): +def _cards(text: str) -> Iterator[tuple[int, int]]: # Cards are written in separated groups, so only whole-group ranges are tried: a card beside # an expiry is still found, while a random digit run gains no Luhn-valid sub-window. for run in re.finditer(r"(? bool: "SI SK SM SO ST SV TL TN TR UA VA VG XK YE".split()) -def _iban_matches(text: str): +def _iban_matches(text: str) -> Iterator[re.Match[str]]: yield from IBAN_PATTERN.finditer(text) for match in LOWER_IBAN_PATTERN.finditer(text): if match.group()[:2].upper() in IBAN_COUNTRIES: yield match -def _ibans(text: str): +def _ibans(text: str) -> Iterator[tuple[int, int]]: for match in _iban_matches(text): candidate = match.group() # The greedy match can swallow following upper-case tokens; drop trailing groups until it checks. @@ -450,7 +460,7 @@ def _is_ipv6(candidate: str) -> bool: return True -def _ipv6s(text: str): +def _ipv6s(text: str) -> Iterator[tuple[int, int]]: for match in IPV6_CANDIDATE.finditer(text): candidate = match.group().rstrip(".") # a sentence's full stop is not part of the address if candidate.endswith(":") and not _is_ipv6(candidate): @@ -469,8 +479,8 @@ def _ipv6s(text: str): r"\.(?:[Xx][Nn]--[A-Za-z0-9-]+|[^\W\d_]{2,})") -def _matches(pattern, accept=None): - def find(text: str): +def _matches(pattern: re.Pattern[str], accept: Callable[[str], bool] | None = None) -> Finder: + def find(text: str) -> Iterator[tuple[int, int]]: for match in pattern.finditer(text): if accept is None or accept(match.group()): yield match.start(), match.end() @@ -479,7 +489,7 @@ def find(text: str): # Precedence for overlapping candidates: checksummed and exact rules before patterns. Every # pattern is anchored on its left so scanning stays linear in the text length. -RULES = ( +RULES: tuple[tuple[str, Finder, float], ...] = ( ("CARD", _cards, 1.0), ("IBAN", _ibans, 1.0), ("EMAIL", _matches(EMAIL_PATTERN), 1.0), @@ -494,9 +504,10 @@ def find(text: str): JOINERS = frozenset({"", "-", ".", "'", "\u2019"}) -def _entities(window, probabilities: list, entities: list, min_score: float, text: str) -> list: +def _entities(window: Any, probabilities: list[list[float]], entities: list[str], min_score: float, + text: str) -> list[dict[str, Any]]: # Each word takes its first word piece's label, the usual BERT aggregation; later pieces only extend it. - words = [] + words: list[list[Any]] = [] for position, word_id in enumerate(window.word_ids): if word_id is None: continue @@ -507,13 +518,15 @@ def _entities(window, probabilities: list, entities: list, min_score: float, tex row = probabilities[position] label_id = max(range(len(row)), key=row.__getitem__) words.append([start, end, NER_LABELS[label_id], float(row[label_id]), word_id]) - pieces, current = [], None + pieces: list[dict[str, Any]] = [] + current: dict[str, Any] | None = None for start, end, label, score, _ in words: prefix, _, kind = label.partition("-") if prefix == "B" or (prefix == "I" and (current is None or current["label"] != kind)): current = {"start": start, "end": end, "label": kind, "scores": [score]} pieces.append(current) elif prefix == "I": + assert current is not None # the branch above took an I- piece without one current["end"] = end current["scores"].append(score) else: @@ -521,7 +534,7 @@ def _entities(window, probabilities: list, entities: list, min_score: float, tex # The model often starts a new entity at the full stop of an initial and is less sure of what follows, so touching # pieces are joined before the threshold and scored by their most confident piece: the initial carries the # surname, and an unsure piece never hides a confident one, as an average over "Jean-Luc Moreau" did. - joined = [] + joined: list[dict[str, Any]] = [] for piece in pieces: score = sum(piece["scores"]) / len(piece["scores"]) if joined and joined[-1]["label"] == piece["label"] and text[joined[-1]["end"]:piece["start"]] in JOINERS: @@ -541,18 +554,19 @@ def _entities(window, probabilities: list, entities: list, min_score: float, tex GIVEN_NAME_AFTER = re.compile(r",\s+([^\W\d_]+)") -def _cased_variant(text: str): +def _cased_variant(text: str) -> str | None: return _recase(text)[0] -def _recase(text: str): +def _recase(text: str) -> tuple[str | None, list[tuple[int, int]]]: """(copy, ranges): a copy of text, offset for offset, with each sentence written all in lower case or all in capitals in title case, or (None, []) when no sentence is. The model was trained on cased news and finds almost no lower-case name; sentence case and mixed case are left alone, since title-casing "the board approved" makes "Board" an organisation, except for a surname in capitals before a comma and a given name, the records convention, which the model otherwise takes for a place or misses. A character whose case change would alter its length stays. The ranges are the runs of rewritten sentences, the only part a second reading needs.""" - chars, start, ranges = list(text), 0, [] + chars, start = list(text), 0 + ranges: list[tuple[int, int]] = [] for end in [match.end() for match in SENTENCE_END.finditer(text)] + [len(text)]: cased = [char for char in text[start:end] if char.isupper() or char.islower()] uncased = bool(cased) and not (any(char.isupper() for char in cased) and any(char.islower() for char in cased)) @@ -575,17 +589,18 @@ def _recase(text: str): return (variant, ranges) if variant != text else (None, []) -def _probabilities(session, names, window) -> list: +def _probabilities(session: Any, names: set[str], window: Any) -> list[list[float]]: import numpy as np inputs = _tensors([window]) outputs = session.run(None, {key: value for key, value in inputs.items() if key in names}) logits = _output(session, outputs, "logits")[0] shifted = np.exp(logits - logits.max(axis=-1, keepdims=True)) - return (shifted / shifted.sum(axis=-1, keepdims=True)).tolist() + rows: list[list[float]] = (shifted / shifted.sum(axis=-1, keepdims=True)).tolist() + return rows -def _ner(text: str, entities: list, min_score: float, model_dir: str, progress) -> list: +def _ner(text: str, entities: list[str], min_score: float, model_dir: str, progress: Progress) -> list[dict[str, Any]]: # Overlapping 512 piece windows cover long text; offsets stay relative to the whole text. The cased model reads the # text, and each run of uncased sentences again from its title-cased copy, and only that run, so a second reading # costs a window or two and cannot relabel a sentence that needed no help. The uncased model then reads the whole @@ -595,13 +610,15 @@ def _ner(text: str, entities: list, min_score: float, model_dir: str, progress) cased = _onnx(Path(model_dir) / "redact", 512, stride=64) uncased = _onnx(Path(model_dir) / "redact-uncased", 512, stride=64) variant, ranges = _recase(text) - readings = [(cased, 0, text)] + [(cased, start, variant[start:end]) for start, end in ranges] + [(uncased, 0, text)] - windows = [] + # _recase returns ranges only with a copy, so the filter drops nothing. + readings = ([(cased, 0, text)] + [(cased, start, variant[start:end]) for start, end in ranges if variant is not None] + + [(uncased, 0, text)]) + windows: list[tuple[Any, set[str], int, str, Any]] = [] for (tokenizer, session, names), offset, reading in readings: encoding = tokenizer.encode(reading) windows.extend((session, names, offset, reading, window) for window in [encoding, *encoding.overflowing]) progress(0, len(windows)) - candidates = [] + candidates: list[dict[str, Any]] = [] for index, (session, names, offset, reading, window) in enumerate(windows): for span in _entities(window, _probabilities(session, names, window), entities, min_score, reading): candidates.append({**span, "start": span["start"] + offset, "end": span["end"] + offset}) @@ -609,11 +626,11 @@ def _ner(text: str, entities: list, min_score: float, model_dir: str, progress) return candidates -def _merge(candidates: list) -> list: +def _merge(candidates: list[dict[str, Any]]) -> list[dict[str, Any]]: # Leftmost first, then longest, then the stronger source, then the more confident: of two model readings of the # same words, a sure "person" beats an unsure "place". A candidate overlapping the kept span extends it instead of # being dropped, so no part of either stays visible. - spans = [] + spans: list[dict[str, Any]] = [] for span in sorted(candidates, key=lambda span: (span["start"], span["start"] - span["end"], span["priority"], -span["score"])): if spans and span["start"] < spans[-1]["end"]: @@ -623,7 +640,7 @@ def _merge(candidates: list) -> list: return spans -def _redact(payload: dict, backend: str, model_dir: str, progress) -> dict: +def _redact(payload: dict[str, Any], backend: str, model_dir: str, progress: Progress) -> dict[str, Any]: text = payload["text"] entities = payload.get("entities", REDACT_DEFAULTS["entities"]) candidates = []