From b6970a27969e924c974dd06704a4cd0dd51b5fb8 Mon Sep 17 00:00:00 2001 From: Lingrui Mei Date: Thu, 2 Jul 2026 22:01:17 +0800 Subject: [PATCH] tito: token-in-token-out session-recording gateway (PR #122 stage B) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Port plugins/tito from origin/abridge/gateway-tito onto master as the second slice of the #122 split: a self-contained workspace member providing the TITO gateway — a FastAPI session server that keeps a token-aligned trajectory per session and proxies chat completions across a pool of OpenAI-compatible sglang replicas, with a native engine (incremental pretokenization, append-only message validation, single-step rollback, mismatch reporting). No vendored Miles; imports nothing from agentix core beyond the namespace. Planned fixes on top of the port: - BackendPool: a downed replica now recovers — report_up is wired into the proxy success path, and a down mark expires after down_cooldown seconds (half-open) so a replica that is never picked isn't blacklisted forever. - BackendPool: sticky session pins are LRU-bounded (max_pins) — the recommended rollout flow never DELETEs its session (the trajectory must survive for harvest), so the pin map grew without bound. - chat route: force stream=false upstream (drop stream_options) — the TITO flow needs the full JSON completion; a stream:true agent previously got a 500 from json-parsing the SSE bytes. Fixes from the pre-PR adversarial review: - chat route: accept tool-call-only assistant turns with content:null (sglang/vLLM emit them routinely) — previously the first tool call of an agentic rollout 502'd. - trust_remote_code is now an explicit opt-in (config field + --trust-remote-code), never a hardcoded True: loading a tokenizer executes Python shipped inside the checkpoint repo. - trajectory: keep only the two reachable checkpoints (single-step rollback) instead of every full prefix+completion token list — O(turns^2) dead memory per session. - chat route: malformed request JSON is a 400; structurally malformed 200 upstream bodies are a clean 502 (not an unhandled 500). - proxy responses pass the upstream body through verbatim — no JSON re-encode (which 500'd on NaN/Infinity and perturbed byte-sensitive bodies). - prefix-mismatch diagnostic no longer raises a bare ValueError when the new sequence is shorter than the stored checkpoint. - CLI: --backend-url is repeatable (multi-replica pool) and --routing-policy is exposed — the pool was Python-API-only. - packaging: requires-python >=3.11 (StrEnum), drop unused setproctitle dep, sglang-specific docs corrected. - concurrency (from the concurrency review pass): a retry chain that outruns the trimmed checkpoint window now re-renders from scratch instead of merging onto an empty prefix (which silently dropped the whole stored history and recorded a corrupt trajectory); rollback is planned+validated before any mutation, so a rejected (400) request no longer commits truncation side effects; the phase-3 interference guard uses a monotonic version counter (num_assistant was ABA-prone); DELETE sets `closing` inside the lock (a cancelled DELETE no longer wedges the session); the request body is read before taking the session lock (a dribbling upload no longer blocks DELETE); the forget-on-delete middleware matches only the exact session resource. Tests: real HTTP-surface coverage (httpx ASGITransport over the actual FastAPI app: session flow, forced non-streaming, tool-call null content, failover + cooldown recovery, error shapes, verbatim passthrough) on top of the ported unit tests. plugins/tito wired into the uv workspace + pyright include/extraPaths. Co-authored-by: FatPigeorz Co-Authored-By: Claude Fable 5 --- plugins/tito/.gitignore | 9 + plugins/tito/LICENSE | 202 +++++++++ plugins/tito/README.md | 110 +++++ plugins/tito/agentix/tito/__init__.py | 21 + plugins/tito/agentix/tito/cli.py | 115 +++++ plugins/tito/agentix/tito/config.py | 98 +++++ plugins/tito/agentix/tito/discovery.py | 98 +++++ plugins/tito/agentix/tito/engine/__init__.py | 7 + plugins/tito/agentix/tito/engine/compare.py | 193 +++++++++ plugins/tito/agentix/tito/engine/errors.py | 33 ++ plugins/tito/agentix/tito/engine/messages.py | 72 ++++ .../tito/agentix/tito/engine/pretokenize.py | 244 +++++++++++ .../tito/agentix/tito/engine/processing.py | 16 + plugins/tito/agentix/tito/engine/render.py | 94 +++++ .../tito/agentix/tito/engine/session_app.py | 231 ++++++++++ .../tito/engine/templates/qwen3_fixed.jinja | 85 ++++ .../tito/agentix/tito/engine/trajectory.py | 275 ++++++++++++ plugins/tito/agentix/tito/gateway.py | 61 +++ plugins/tito/agentix/tito/pool.py | 119 ++++++ plugins/tito/agentix/tito/server.py | 118 ++++++ plugins/tito/agentix/tito/tokenizer.py | 29 ++ plugins/tito/pyproject.toml | 69 +++ plugins/tito/tests/package/test_cli.py | 77 ++++ .../tests/package/test_config_discovery.py | 137 ++++++ plugins/tito/tests/package/test_engine.py | 294 +++++++++++++ .../tito/tests/package/test_import_surface.py | 40 ++ plugins/tito/tests/test_gateway_http.py | 393 ++++++++++++++++++ plugins/tito/tests/test_pool.py | 105 +++++ plugins/tito/tests/test_pool_routing.py | 117 ++++++ pyproject.toml | 3 + uv.lock | 119 +++++- 31 files changed, 3566 insertions(+), 18 deletions(-) create mode 100644 plugins/tito/.gitignore create mode 100644 plugins/tito/LICENSE create mode 100644 plugins/tito/README.md create mode 100644 plugins/tito/agentix/tito/__init__.py create mode 100644 plugins/tito/agentix/tito/cli.py create mode 100644 plugins/tito/agentix/tito/config.py create mode 100644 plugins/tito/agentix/tito/discovery.py create mode 100644 plugins/tito/agentix/tito/engine/__init__.py create mode 100644 plugins/tito/agentix/tito/engine/compare.py create mode 100644 plugins/tito/agentix/tito/engine/errors.py create mode 100644 plugins/tito/agentix/tito/engine/messages.py create mode 100644 plugins/tito/agentix/tito/engine/pretokenize.py create mode 100644 plugins/tito/agentix/tito/engine/processing.py create mode 100644 plugins/tito/agentix/tito/engine/render.py create mode 100644 plugins/tito/agentix/tito/engine/session_app.py create mode 100644 plugins/tito/agentix/tito/engine/templates/qwen3_fixed.jinja create mode 100644 plugins/tito/agentix/tito/engine/trajectory.py create mode 100644 plugins/tito/agentix/tito/gateway.py create mode 100644 plugins/tito/agentix/tito/pool.py create mode 100644 plugins/tito/agentix/tito/server.py create mode 100644 plugins/tito/agentix/tito/tokenizer.py create mode 100644 plugins/tito/pyproject.toml create mode 100644 plugins/tito/tests/package/test_cli.py create mode 100644 plugins/tito/tests/package/test_config_discovery.py create mode 100644 plugins/tito/tests/package/test_engine.py create mode 100644 plugins/tito/tests/package/test_import_surface.py create mode 100644 plugins/tito/tests/test_gateway_http.py create mode 100644 plugins/tito/tests/test_pool.py create mode 100644 plugins/tito/tests/test_pool_routing.py diff --git a/plugins/tito/.gitignore b/plugins/tito/.gitignore new file mode 100644 index 0000000..9d78d95 --- /dev/null +++ b/plugins/tito/.gitignore @@ -0,0 +1,9 @@ +__pycache__/ +*.py[cod] +.humanize/ +.pytest_cache/ +.ruff_cache/ +.mypy_cache/ +dist/ +build/ +*.egg-info/ diff --git a/plugins/tito/LICENSE b/plugins/tito/LICENSE new file mode 100644 index 0000000..d645695 --- /dev/null +++ b/plugins/tito/LICENSE @@ -0,0 +1,202 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/plugins/tito/README.md b/plugins/tito/README.md new file mode 100644 index 0000000..f1076d4 --- /dev/null +++ b/plugins/tito/README.md @@ -0,0 +1,110 @@ +# agentix-tito — TITO Gateway + +An Agentix plugin (`import agentix.tito`) that records **token-aligned** +agent↔model trajectories. It sits between an agent and an OpenAI-compatible +inference backend (e.g. sglang) as a session-scoped proxy, and accumulates the +exact `input_ids` / completion token IDs of every turn — the trajectory an RL +trainer needs, with no host-side re-tokenization. + +This is a **native implementation** of the TITO token-alignment engine +(`agentix.tito.engine`): no vendored training-framework code and no `sglang` +dependency. The engine tokenizes prompts itself with `transformers` + +`tokenizers` + a fixed Jinja chat template. + +## TITO in one paragraph + +TITO = *token-in, token-out*. Instead of re-tokenizing the rendered chat +transcript host-side (which can drift from what the model actually saw), the +gateway: + +1. **pretokenizes** each request's messages to `input_ids` and sends those to + the backend (token-in), +2. reads the **exact completion token IDs** back from + `meta_info.output_token_logprobs` (token-out), +3. reuses the **byte-identical token prefix** across turns, tokenizing only the + newly-appended non-assistant messages (tool/user/system) as a suffix in a + synthetic context, and +4. on read, **audits** the accumulated trajectory against a from-scratch render + (`compute_session_mismatch`) so any tokenizer drift is detected, not hidden. + +The algorithm is **model-agnostic** (base `TITOTokenizer`); a model family is a +fixed chat template plus a tiny boundary fixup — e.g. `Qwen3TITOTokenizer` +re-inserts the `\n` after `<|im_end|>` that the model omits when it stops. + +## Install + +It is a member of the Agentix uv workspace, installed editable with the rest: + +```bash +uv sync --all-packages --all-extras +``` + +Its runtime deps (`transformers`, `tokenizers`, `jinja2`, …) are isolated to +this plugin — agentix core and other plugins never pull them. + +## CLI + +```bash +agentix-tito serve \ + --hf-checkpoint Qwen/Qwen3-4B \ + --backend-url http://127.0.0.1:30000 \ + --tito-model qwen3 \ + --session-server-port 30001 +``` + +`--tito-model` selects the tokenizer family (`qwen3`, or `default` for the +tokenizer's own template). `--backend-url` may be omitted to auto-discover a +local backend (see `agentix.tito.discovery`). Run `agentix-tito serve -h` for +the full list. + +## HTTP surface + +- `POST /sessions` → `{session_id}` +- `POST /sessions/{id}/v1/chat/completions` — proxied chat completion; the + gateway forces `logprobs`/`return_meta_info`, injects the pretokenized + `input_ids`, and appends a token-aligned checkpoint. +- `GET /sessions/{id}` — records + metadata, incl. `accumulated_token_ids` and + `tito_session_mismatch` (empty list ⇒ byte-identical to a fresh render). +- `DELETE /sessions/{id}` — close the session and forget its pool pin. + +Multiple backend replicas are supported via `BackendPool`: requests are pinned +sticky-by-`session_id` for prefix-cache locality, and a replica is marked down +on a transport error. + +## Python API + +```python +from agentix.tito import TITOGateway, TITOGatewayConfig + +TITOGateway(TITOGatewayConfig( + hf_checkpoint="Qwen/Qwen3-4B", + backend_url="http://127.0.0.1:30000", + tito_model="qwen3", +)).run() +``` + +`agentix.tito.get_tito_tokenizer(tokenizer, "qwen3")` builds the engine +tokenizer directly if you only want incremental pretokenization. + +## Layout + +```text +agentix/tito/ +├── gateway.py / server.py / pool.py / discovery.py / config.py / cli.py +└── engine/ — the native TITO token-alignment engine + ├── pretokenize.py — TITOTokenizer (+ Qwen3TITOTokenizer) + ├── compare.py — special-token-segment mismatch audit + ├── trajectory.py — LinearTrajectory + SessionRegistry + ├── session_app.py — FastAPI session routes + ├── messages.py / render.py / processing.py / errors.py + └── templates/qwen3_fixed.jinja +``` + +## Tests + +```bash +pytest plugins/tito/tests +``` + +The engine tests are self-contained — they build a tiny in-memory tokenizer, so +no model download or GPU is required. diff --git a/plugins/tito/agentix/tito/__init__.py b/plugins/tito/agentix/tito/__init__.py new file mode 100644 index 0000000..1c76b06 --- /dev/null +++ b/plugins/tito/agentix/tito/__init__.py @@ -0,0 +1,21 @@ +"""Agentix TITO plugin — token-in-token-out session-recording gateway. + +A native reimplementation of the TITO token-alignment engine (see +`agentix.tito.engine`); no vendored training-framework code and no sglang +dependency. +""" + +from .config import TITOGatewayConfig +from .discovery import discover_backend_url +from .gateway import TITOGateway +from .server import SessionServer +from .tokenizer import TITOTokenizerType, get_tito_tokenizer + +__all__ = [ + "TITOGateway", + "TITOGatewayConfig", + "SessionServer", + "TITOTokenizerType", + "discover_backend_url", + "get_tito_tokenizer", +] diff --git a/plugins/tito/agentix/tito/cli.py b/plugins/tito/agentix/tito/cli.py new file mode 100644 index 0000000..1eb1f3a --- /dev/null +++ b/plugins/tito/agentix/tito/cli.py @@ -0,0 +1,115 @@ +"""Command-line entrypoint for the Agentix TITO gateway.""" + +from __future__ import annotations + +import argparse +import sys + +from .config import TITOGatewayConfig +from .gateway import TITOGateway +from .tokenizer import TITOTokenizerType + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="agentix-tito", + description="Agentix TITO gateway — token-in-token-out session-recording proxy.", + ) + subparsers = parser.add_subparsers(dest="command") + _add_serve_parser(subparsers) + return parser + + +def _add_serve_parser(subparsers: argparse._SubParsersAction) -> argparse.ArgumentParser: + serve = subparsers.add_parser("serve", help="Start the TITO gateway server.") + _add_serve_arguments(serve) + return serve + + +def _add_serve_arguments(parser: argparse.ArgumentParser) -> None: + parser.add_argument("--hf-checkpoint", required=True, help="HuggingFace model ID or local checkpoint path.") + parser.add_argument( + "--backend-url", + action="append", + default=None, + help="OpenAI-compatible backend URL to proxy to; repeat for a multi-replica pool.", + ) + parser.add_argument( + "--routing-policy", + choices=["sticky", "round_robin"], + default="sticky", + help="Multi-replica routing: pin each session to one replica (sticky) or spread every request.", + ) + parser.add_argument( + "--trust-remote-code", + action="store_true", + help="Allow the tokenizer load to execute Python shipped inside the checkpoint repo (off by default).", + ) + parser.add_argument("--chat-template-path", default=None, help="Optional fixed chat template path.") + parser.add_argument( + "--tito-model", + choices=[item.value for item in TITOTokenizerType], + default=TITOTokenizerType.DEFAULT.value, + help="TITO tokenizer family (qwen3, or default for the tokenizer's own template).", + ) + parser.add_argument( + "--tito-allowed-append-roles", + nargs="+", + choices=["tool", "user", "system"], + default=["tool"], + help="Roles allowed after an assistant turn; tool is the default.", + ) + parser.add_argument("--session-server-ip", default="127.0.0.1", help="Gateway bind host.") + parser.add_argument("--session-server-port", type=int, default=30000, help="Gateway bind port.") + parser.add_argument("--router-timeout", type=float, default=600.0, help="Proxy timeout in seconds.") + parser.add_argument( + "--backend-probe-candidate", + action="append", + default=None, + metavar="URL", + help="Local backend URL candidate to probe after explicit and environment URLs; repeatable.", + ) + parser.add_argument( + "--backend-probe-timeout", + type=float, + default=0.25, + help="Per-endpoint backend probe timeout in seconds.", + ) + + +def _serve(args: argparse.Namespace) -> int: + urls: list[str] = args.backend_url or [] + config = TITOGatewayConfig.from_cli_values( + hf_checkpoint=args.hf_checkpoint, + backend_url=urls[0] if len(urls) == 1 else None, + backend_urls=urls if len(urls) > 1 else None, + routing_policy=args.routing_policy, + trust_remote_code=args.trust_remote_code, + chat_template_path=args.chat_template_path, + tito_model=args.tito_model, + tito_allowed_append_roles=args.tito_allowed_append_roles, + session_server_ip=args.session_server_ip, + session_server_port=args.session_server_port, + router_timeout=args.router_timeout, + backend_probe_candidates=args.backend_probe_candidate, + backend_probe_timeout=args.backend_probe_timeout, + ) + TITOGateway(config).run() + return 0 + + +def main(argv: list[str] | None = None) -> int: + parser = build_parser() + raw_args = list(sys.argv[1:] if argv is None else argv) + if not raw_args or raw_args[0] not in {"serve", "-h", "--help"}: + raw_args.insert(0, "serve") + args = parser.parse_args(raw_args) + try: + return _serve(args) + except Exception as exc: # noqa: BLE001 + print(f"agentix-tito: error: {exc}", file=sys.stderr) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/plugins/tito/agentix/tito/config.py b/plugins/tito/agentix/tito/config.py new file mode 100644 index 0000000..62c8ea4 --- /dev/null +++ b/plugins/tito/agentix/tito/config.py @@ -0,0 +1,98 @@ +"""Configuration objects for the TITO Gateway wrapper.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from .discovery import DEFAULT_BACKEND_PROBE_CANDIDATES + +_VALID_APPEND_ROLES = frozenset({"tool", "user", "system"}) + + +@dataclass(frozen=True) +class TITOGatewayConfig: + """Configuration for the standalone TITO gateway wrapper.""" + + hf_checkpoint: str + backend_url: str | None = None + # Explicit multi-backend pool (sglang replicas — the token-recording chat + # flow needs sglang's meta_info extension). When set, these are used as-is + # and single-URL discovery is skipped; `backend_url` is left as the first + # entry for callers that read it. + backend_urls: tuple[str, ...] = () + routing_policy: str = "sticky" + # Execute Python shipped inside the checkpoint repo when loading the + # tokenizer. Off by default; opt in only for checkpoints you trust. + trust_remote_code: bool = False + chat_template_path: str | None = None + tito_model: str = "default" + tito_allowed_append_roles: tuple[str, ...] = ("tool",) + session_server_ip: str = "127.0.0.1" + session_server_port: int = 30000 + router_timeout: float = 600.0 + backend_probe_candidates: tuple[str, ...] = field(default_factory=lambda: DEFAULT_BACKEND_PROBE_CANDIDATES) + backend_probe_timeout: float = 0.25 + + def __post_init__(self) -> None: + if not self.hf_checkpoint: + raise ValueError("hf_checkpoint is required for TITO token tracking") + + normalized_roles = tuple(dict.fromkeys(role.lower() for role in self.tito_allowed_append_roles)) + invalid = sorted(set(normalized_roles) - _VALID_APPEND_ROLES) + if invalid: + raise ValueError(f"unsupported tito append roles: {invalid}") + object.__setattr__(self, "tito_allowed_append_roles", normalized_roles or ("tool",)) + + if self.routing_policy not in ("sticky", "round_robin"): + raise ValueError( + f"routing_policy must be 'sticky' or 'round_robin'; got {self.routing_policy!r}" + ) + + @classmethod + def from_cli_values( + cls, + *, + hf_checkpoint: str, + backend_url: str | None, + chat_template_path: str | None, + tito_model: str, + tito_allowed_append_roles: list[str], + session_server_ip: str, + session_server_port: int, + router_timeout: float, + backend_urls: list[str] | None = None, + routing_policy: str = "sticky", + trust_remote_code: bool = False, + backend_probe_candidates: list[str] | None = None, + backend_probe_timeout: float = 0.25, + ) -> TITOGatewayConfig: + return cls( + hf_checkpoint=hf_checkpoint, + backend_url=backend_url, + backend_urls=tuple(backend_urls or ()), + routing_policy=routing_policy, + trust_remote_code=trust_remote_code, + chat_template_path=chat_template_path, + tito_model=tito_model, + tito_allowed_append_roles=tuple(tito_allowed_append_roles), + session_server_ip=session_server_ip, + session_server_port=session_server_port, + router_timeout=router_timeout, + backend_probe_candidates=tuple(backend_probe_candidates or DEFAULT_BACKEND_PROBE_CANDIDATES), + backend_probe_timeout=backend_probe_timeout, + ) + + def as_session_args(self): + """Return an argparse-like namespace consumed by the engine session routes.""" + from types import SimpleNamespace + + return SimpleNamespace( + hf_checkpoint=self.hf_checkpoint, + chat_template_path=self.chat_template_path, + tito_model=self.tito_model, + tito_allowed_append_roles=list(self.tito_allowed_append_roles), + trust_remote_code=self.trust_remote_code, + session_server_ip=self.session_server_ip, + session_server_port=self.session_server_port, + router_timeout=self.router_timeout, + ) diff --git a/plugins/tito/agentix/tito/discovery.py b/plugins/tito/agentix/tito/discovery.py new file mode 100644 index 0000000..afc8b69 --- /dev/null +++ b/plugins/tito/agentix/tito/discovery.py @@ -0,0 +1,98 @@ +"""Backend URL discovery for wrapping an OpenAI-compatible server.""" + +from __future__ import annotations + +import logging +import os +from collections.abc import Callable, Iterable, Mapping +from urllib.error import HTTPError, URLError +from urllib.request import Request, urlopen + +DEFAULT_BACKEND_ENV_VARS = ("TITO_BACKEND_URL", "OPENAI_BASE_URL", "SGLANG_BASE_URL") +DEFAULT_BACKEND_PROBE_CANDIDATES = ( + "http://127.0.0.1:8000", + "http://localhost:8000", + "http://127.0.0.1:30000", + "http://localhost:30000", +) +BACKEND_PROBE_PATHS = ("/health", "/v1/models") + +logger = logging.getLogger(__name__) + + +def normalize_backend_url(url: str) -> str: + normalized = url.strip().rstrip("/") + if not normalized: + raise ValueError("backend URL is empty") + if "://" not in normalized: + normalized = f"http://{normalized}" + return normalized + + +def _probe_endpoint(url: str, timeout: float) -> bool: + request = Request(url, method="GET") + try: + with urlopen(request, timeout=timeout) as response: + return 200 <= response.status < 300 + except HTTPError as exc: + return 200 <= exc.code < 300 + except (OSError, URLError, TimeoutError, ValueError): + return False + + +def probe_backend_url( + candidate_url: str, + *, + timeout: float = 0.25, + endpoint_probe: Callable[[str, float], bool] | None = None, +) -> str | None: + """Return the successful probe path for a backend candidate, if live.""" + backend_url = normalize_backend_url(candidate_url) + probe = _probe_endpoint if endpoint_probe is None else endpoint_probe + for path in BACKEND_PROBE_PATHS: + if probe(f"{backend_url}{path}", timeout): + return path + return None + + +def discover_backend_url( + explicit_url: str | None = None, + *, + env: Mapping[str, str] | None = None, + env_vars: Iterable[str] = DEFAULT_BACKEND_ENV_VARS, + probe_candidates: Iterable[str] | None = DEFAULT_BACKEND_PROBE_CANDIDATES, + probe_timeout: float = 0.25, +) -> str: + """Resolve the backend URL with deterministic precedence. + + Explicit config wins, followed by environment variables in + ``DEFAULT_BACKEND_ENV_VARS`` order, followed by configured local probe + candidates in the supplied order. + """ + if explicit_url: + backend_url = normalize_backend_url(explicit_url) + logger.info("Selected backend URL from explicit config: %s", backend_url) + return backend_url + + source = os.environ if env is None else env + for key in env_vars: + value = source.get(key) + if value: + backend_url = normalize_backend_url(value) + logger.info("Selected backend URL from %s: %s", key, backend_url) + return backend_url + + candidates = tuple(probe_candidates or ()) + for candidate in candidates: + backend_url = normalize_backend_url(candidate) + live_path = probe_backend_url(backend_url, timeout=probe_timeout) + if live_path: + logger.info("Selected backend URL from probe %s via %s", backend_url, live_path) + return backend_url + + names = ", ".join(env_vars) + candidate_text = ", ".join(candidates) if candidates else "none configured" + raise RuntimeError( + "backend URL not found; pass --backend-url, " + f"set one of: {names}, or start a live backend on one of: {candidate_text}" + ) diff --git a/plugins/tito/agentix/tito/engine/__init__.py b/plugins/tito/agentix/tito/engine/__init__.py new file mode 100644 index 0000000..1c69946 --- /dev/null +++ b/plugins/tito/agentix/tito/engine/__init__.py @@ -0,0 +1,7 @@ +"""Agentix's native TITO engine — token-in token-out pretokenization, session +trajectory, and mismatch-audit logic. + +Model-agnostic core lives here; per-model behavior is a small amount of data +(a fixed chat template) plus a couple of constants and an optional boundary +fixup. See `pretokenize.TITOTokenizer` for the algorithm. +""" diff --git a/plugins/tito/agentix/tito/engine/compare.py b/plugins/tito/agentix/tito/engine/compare.py new file mode 100644 index 0000000..7981a58 --- /dev/null +++ b/plugins/tito/agentix/tito/engine/compare.py @@ -0,0 +1,193 @@ +"""Token-sequence comparator: segment by special tokens, classify mismatches. + +Used to check that an incrementally-accumulated trajectory tokenizes identically +to a from-scratch render. The comparison is structural: the special-token skeleton +and non-assistant content must match exactly; assistant content may differ (the +model's own tokens) and is reported as a soft mismatch. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import StrEnum +from typing import Any + + +class MismatchType(StrEnum): + # Segment count or special/content pattern differs — structural break. + SPECIAL_TOKEN_COUNT = "special_token_count" + # Aligned special-token segment holds a different special token. + SPECIAL_TOKEN_TYPE = "special_token_type" + # Non-assistant content (system/user/tool) differs — the prompt drifted. + NON_ASSISTANT_TEXT = "non_assistant_text" + # Assistant content differs — expected and non-severe (model's own tokens). + ASSISTANT_TEXT = "assistant_text" + + +@dataclass +class Segment: + token_ids: list[int] + is_special: bool = False + + +@dataclass +class Mismatch: + type: MismatchType + segment_index: int + expected_text: str = "" + actual_text: str = "" + detail: str = "" + + def to_dict(self) -> dict[str, Any]: + return { + "type": self.type.value, + "segment_index": self.segment_index, + "expected_text": self.expected_text, + "actual_text": self.actual_text, + "detail": self.detail, + } + + +class TokenSeqComparator: + """Segment two token-ID sequences at special-token boundaries and compare. + + `assistant_start_str` (e.g. ``"<|im_start|>assistant"``) classifies a content + segment as assistant vs non-assistant. `special_token_ids`, if given, overrides + the set collected from the tokenizer. `trim_trailing_ids` are stripped from both + tails before comparison (a stop token the model emits but the template doesn't). + """ + + def __init__( + self, + tokenizer: Any, + *, + assistant_start_str: str | None, + special_token_ids: set[int] | None = None, + trim_trailing_ids: frozenset[int] | set[int] | None = None, + ) -> None: + self.tokenizer = tokenizer + self._assistant_start_str = assistant_start_str + self._special_ids = ( + set(special_token_ids) if special_token_ids is not None else self.collect_special_ids(tokenizer) + ) + self._trim_trailing_ids = set(trim_trailing_ids) if trim_trailing_ids else None + + @staticmethod + def collect_special_ids(tokenizer: Any) -> set[int]: + """Token IDs flagged ``special=True`` by the tokenizer. Content tokens a role + produces (e.g. ````) are NOT special, so they aren't collected here.""" + ids: set[int] = set(getattr(tokenizer, "all_special_ids", []) or []) + decoder = getattr(tokenizer, "added_tokens_decoder", None) + if decoder: + ids |= {k for k, v in decoder.items() if getattr(v, "special", False)} + return ids + + def segment_by_special_tokens(self, token_ids: list[int]) -> list[Segment]: + """Each special token is its own single-ID segment; consecutive non-special + tokens group into one content segment.""" + segments: list[Segment] = [] + current: list[int] = [] + for tid in token_ids: + if tid in self._special_ids: + if current: + segments.append(Segment(token_ids=current)) + current = [] + segments.append(Segment(token_ids=[tid], is_special=True)) + else: + current.append(tid) + if current: + segments.append(Segment(token_ids=current)) + return segments + + def compare_sequences( + self, + expected_ids: list[int], + actual_ids: list[int], + trim_trailing_ids: frozenset[int] | set[int] | None = None, + ) -> list[Mismatch]: + trim = self._trim_trailing_ids or set() + if trim_trailing_ids: + trim = trim | trim_trailing_ids + if trim: + expected_ids = _trim_trailing(expected_ids, trim) + actual_ids = _trim_trailing(actual_ids, trim) + + exp_segs = self.segment_by_special_tokens(expected_ids) + act_segs = self.segment_by_special_tokens(actual_ids) + + structural = self._check_segment_structure(exp_segs, act_segs) + if structural is not None: + return [structural] + + mismatches: list[Mismatch] = [] + for idx, (exp, act) in enumerate(zip(exp_segs, act_segs, strict=True)): + is_assistant = self._is_assistant_content(exp_segs, idx) and self._is_assistant_content(act_segs, idx) + m = self._compare_single_segment(idx, exp, act, is_assistant_content=is_assistant) + if m is not None: + mismatches.append(m) + return mismatches + + def _check_segment_structure(self, exp_segs: list[Segment], act_segs: list[Segment]) -> Mismatch | None: + if len(exp_segs) != len(act_segs): + detail = f"segment count differs: expected {len(exp_segs)}, got {len(act_segs)}" + elif [s.is_special for s in exp_segs] != [s.is_special for s in act_segs]: + detail = "segment structure (special/content pattern) differs" + else: + return None + return Mismatch( + type=MismatchType.SPECIAL_TOKEN_COUNT, + segment_index=-1, + expected_text=self._describe_structure(exp_segs), + actual_text=self._describe_structure(act_segs), + detail=detail, + ) + + def _compare_single_segment( + self, idx: int, exp: Segment, act: Segment, *, is_assistant_content: bool + ) -> Mismatch | None: + if exp.is_special: + if exp.token_ids != act.token_ids: + return Mismatch( + type=MismatchType.SPECIAL_TOKEN_TYPE, + segment_index=idx, + expected_text=self._decode(exp.token_ids), + actual_text=self._decode(act.token_ids), + ) + return None + exp_text = self._decode(exp.token_ids) + act_text = self._decode(act.token_ids) + if exp_text == act_text: + return None + return Mismatch( + type=MismatchType.ASSISTANT_TEXT if is_assistant_content else MismatchType.NON_ASSISTANT_TEXT, + segment_index=idx, + expected_text=exp_text, + actual_text=act_text, + ) + + def _is_assistant_content(self, segments: list[Segment], idx: int) -> bool: + if self._assistant_start_str is None: + return False + if segments[idx].is_special or idx == 0: + return False + prev = segments[idx - 1] + if not prev.is_special: + return False + special_text = self._decode(prev.token_ids) + content_prefix = self._decode(segments[idx].token_ids[:20]) + return (special_text + content_prefix).startswith(self._assistant_start_str) + + def _decode(self, token_ids: list[int]) -> str: + return self.tokenizer.decode(token_ids, skip_special_tokens=False) + + def _describe_structure(self, segments: list[Segment]) -> str: + return " ".join( + f"[{self._decode(s.token_ids)}]" if s.is_special else f"({len(s.token_ids)} tokens)" for s in segments + ) + + +def _trim_trailing(ids: list[int], to_remove: set[int]) -> list[int]: + end = len(ids) + while end > 0 and ids[end - 1] in to_remove: + end -= 1 + return ids[:end] diff --git a/plugins/tito/agentix/tito/engine/errors.py b/plugins/tito/agentix/tito/engine/errors.py new file mode 100644 index 0000000..d93aaa8 --- /dev/null +++ b/plugins/tito/agentix/tito/engine/errors.py @@ -0,0 +1,33 @@ +"""Session error hierarchy. Each carries the HTTP status the gateway returns.""" + +from __future__ import annotations + + +class SessionError(Exception): + """Base class for all session-related errors.""" + + status_code: int = 500 + + +class SessionNotFoundError(SessionError): + """The requested session ID does not exist.""" + + status_code: int = 404 + + +class MessageValidationError(SessionError): + """Request messages aren't a valid append-only extension (or a rollback failed).""" + + status_code: int = 400 + + +class TokenizationError(SessionError): + """A TITO tokenization invariant was violated (e.g. pretokenized prefix mismatch).""" + + status_code: int = 500 + + +class UpstreamResponseError(SessionError): + """The upstream sglang response is invalid or unexpected (missing meta_info, etc.).""" + + status_code: int = 502 diff --git a/plugins/tito/agentix/tito/engine/messages.py b/plugins/tito/agentix/tito/engine/messages.py new file mode 100644 index 0000000..2035dfb --- /dev/null +++ b/plugins/tito/agentix/tito/engine/messages.py @@ -0,0 +1,72 @@ +"""Message-level helpers used by the session state machine. + +`message_matches` compares only the fields that affect chat-template tokenization, +so a stored assistant message and a resent one are considered equal iff they +tokenize identically. `assert_messages_append_only_with_allowed_role` enforces that +each turn extends the stored history without rewriting it. +""" + +from __future__ import annotations + +from typing import Any + +# Keys a chat template actually reads. Extra client-injected keys +# (provider_specific_fields, etc.) don't affect tokenization, so we ignore them. +TEMPLATE_RELEVANT_KEYS = ("role", "content", "reasoning_content", "tool_calls") + +DEFAULT_APPEND_ROLES: list[str] = ["tool"] + + +def normalize_value(value: Any) -> Any: + """Collapse the falsy sentinels that render identically in Jinja2 (None, "", []) + to None. Non-falsy content — including whitespace like trailing newlines — is + returned as-is, because boundary characters must tokenize identically.""" + if value is None or value == "" or value == []: + return None + return value + + +def message_matches(stored: dict[str, Any], new: dict[str, Any]) -> bool: + for key in TEMPLATE_RELEVANT_KEYS: + if normalize_value(stored.get(key)) != normalize_value(new.get(key)): + return False + return True + + +def assert_messages_append_only_with_allowed_role( + stored_messages: list[dict[str, Any]], + new_messages: list[dict[str, Any]], + allowed_append_roles: list[str] = DEFAULT_APPEND_ROLES, +) -> None: + """Assert *new_messages* is an append-only extension of *stored_messages*: the + stored prefix matches (by template-relevant keys) and each appended message's + role is in *allowed_append_roles*. Raises ValueError otherwise.""" + if not stored_messages: + return + + if len(new_messages) < len(stored_messages): + raise ValueError( + f"new messages ({len(new_messages)}) are fewer than stored messages ({len(stored_messages)})", + new_messages, + stored_messages, + ) + + for i, stored_msg in enumerate(stored_messages): + if not message_matches(stored_msg, new_messages[i]): + diffs = { + key: {"stored": repr(stored_msg.get(key))[:200], "new": repr(new_messages[i].get(key))[:200]} + for key in TEMPLATE_RELEVANT_KEYS + if stored_msg.get(key) != new_messages[i].get(key) + } + raise ValueError( + f"message mismatch at index {i} " + f"(role: stored={stored_msg.get('role')}, new={new_messages[i].get('role')}). " + f"Diffs: {diffs}" + ) + + for j, msg in enumerate(new_messages[len(stored_messages):]): + if msg.get("role") not in allowed_append_roles: + raise ValueError( + f"appended message at index {len(stored_messages) + j} " + f"has role={msg.get('role')!r}, allowed={allowed_append_roles}" + ) diff --git a/plugins/tito/agentix/tito/engine/pretokenize.py b/plugins/tito/agentix/tito/engine/pretokenize.py new file mode 100644 index 0000000..4d22aa0 --- /dev/null +++ b/plugins/tito/agentix/tito/engine/pretokenize.py @@ -0,0 +1,244 @@ +"""TITO tokenizer — incremental tokenization for pretokenized-prefix reuse. + +The base `TITOTokenizer` holds the whole model-agnostic algorithm: it computes the +token IDs for non-assistant messages (tool/user/system) appended after the +assistant's generated tokens, by rendering each segment in a minimal synthetic +context and taking the suffix, then merges them onto the stored prefix. A model +subclass only fixes boundary tokens at the junction (e.g. Qwen3's missing newline +after `<|im_end|>`) and points at its fixed chat template. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from .compare import TokenSeqComparator +from .messages import assert_messages_append_only_with_allowed_role +from .render import apply_chat_template + +TEMPLATE_DIR = Path(__file__).parent / "templates" +_VALID_ROLES = frozenset({"tool", "user", "system"}) +_DUMMY_SYSTEM: dict[str, Any] = {"role": "system", "content": "dummy system"} + + +def _build_dummy_assistant(tool_responses: list[dict[str, Any]]) -> dict[str, Any]: + """A dummy assistant whose tool_calls match *tool_responses*, so the template + renders the following tool-response turn boundaries correctly.""" + return { + "role": "assistant", + "content": "", + "reasoning_content": " ", + "tool_calls": [ + { + "id": resp.get("tool_call_id") or f"call0000{i}", + "type": "function", + "function": {"name": resp.get("name") or "dummy_func", "arguments": {}}, + } + for i, resp in enumerate(tool_responses) + ], + } + + +class TITOTokenizer: + """Incremental tokenization + prefix merging for appended non-assistant turns.""" + + max_trim_tokens: int = 0 + trailing_token_ids: frozenset[int] = frozenset() + reasoning_parser: str | None = None + tool_call_parser: str | None = None + + def __init__( + self, + tokenizer: Any, + chat_template_kwargs: dict[str, Any] | None = None, + assistant_start_str: str | None = None, + special_token_ids: set[int] | None = None, + allowed_append_roles: list[str] | None = None, + ) -> None: + self.tokenizer = tokenizer + self.chat_template_kwargs = chat_template_kwargs or {} + self._assistant_start_str = assistant_start_str + self.allowed_append_roles: list[str] = allowed_append_roles if allowed_append_roles is not None else ["tool"] + self.special_token_ids = special_token_ids + + def create_comparator(self) -> TokenSeqComparator: + return TokenSeqComparator( + self.tokenizer, + assistant_start_str=self._assistant_start_str, + special_token_ids=self.special_token_ids, + trim_trailing_ids=self.trailing_token_ids or None, + ) + + def render_messages( + self, + messages: list[dict[str, Any]], + *, + add_generation_prompt: bool, + tools: list[dict[str, Any]] | None = None, + tokenize: bool = False, + ) -> Any: + return apply_chat_template( + messages, + tokenizer=self.tokenizer, + tokenize=tokenize, + add_generation_prompt=add_generation_prompt, + tools=tools, + **self.chat_template_kwargs, + ) + + def _encode_text(self, text: str) -> list[int]: + return self.tokenizer.encode(text, add_special_tokens=False) + + def _split_appended_segments(self, appended_messages: list[dict[str, Any]]) -> list[list[dict[str, Any]]]: + segments: list[list[dict[str, Any]]] = [] + i = 0 + while i < len(appended_messages): + role = appended_messages[i]["role"] + if role == "tool": + j = i + 1 + while j < len(appended_messages) and appended_messages[j]["role"] == "tool": + j += 1 + segments.append(appended_messages[i:j]) + i = j + continue + if role in {"user", "system"}: + segments.append([appended_messages[i]]) + i += 1 + continue + raise ValueError(f"unsupported appended role for TITO segmentation: {role}") + return segments + + def _tokenize_rendered_suffix( + self, + base_messages: list[dict[str, Any]], + appended_messages: list[dict[str, Any]], + *, + tools: list[dict[str, Any]] | None = None, + add_generation_prompt: bool = False, + ) -> list[int]: + text_without = self.render_messages(base_messages, add_generation_prompt=False, tools=tools) + text_with = self.render_messages( + base_messages + appended_messages, add_generation_prompt=add_generation_prompt, tools=tools + ) + if not text_with.startswith(text_without): + roles = [m["role"] for m in appended_messages] if appended_messages else ["generation_prompt"] + raise ValueError(f"rendered suffix diff failed for {roles}") + return self._encode_text(text_with[len(text_without):]) + + def _tokenize_tool_segment( + self, appended_messages: list[dict[str, Any]], tools: list[dict[str, Any]] | None = None + ) -> list[int]: + return self._tokenize_rendered_suffix( + [_DUMMY_SYSTEM, _build_dummy_assistant(appended_messages)], appended_messages, tools=tools + ) + + def _tokenize_user_and_system_segment( + self, appended_message: dict[str, Any], tools: list[dict[str, Any]] | None = None + ) -> list[int]: + return self._tokenize_rendered_suffix([_DUMMY_SYSTEM], [appended_message], tools=tools) + + def tokenize_additional_non_assistant( + self, + old_messages: list[dict[str, Any]], + new_messages: list[dict[str, Any]], + tools: list[dict[str, Any]] | None = None, + ) -> list[int]: + """Incremental token IDs (incl. the next generation prompt) for the + non-assistant messages appended after the pretokenized prefix.""" + assert_messages_append_only_with_allowed_role(old_messages, new_messages, self.allowed_append_roles) + appended_messages = new_messages[len(old_messages):] + incremental: list[int] = [] + for segment in self._split_appended_segments(appended_messages): + role = segment[0]["role"] + if role == "tool": + incremental.extend(self._tokenize_tool_segment(segment, tools)) + elif role in ("user", "system"): + incremental.extend(self._tokenize_user_and_system_segment(segment[0], tools)) + else: + raise ValueError(f"unsupported appended role for TITO tokenization: {role}") + return incremental + self._tokenize_rendered_suffix( + new_messages, [], tools=tools, add_generation_prompt=True + ) + + def merge_tokens( + self, + old_messages: list[dict[str, Any]], + new_messages: list[dict[str, Any]], + pretokenized_token_ids: list[int], + tools: list[dict[str, Any]] | None = None, + ) -> list[int]: + """Default: concatenate the stored prefix with the incremental tokens.""" + incremental = self.tokenize_additional_non_assistant(old_messages, new_messages, tools) + return list(pretokenized_token_ids) + incremental + + +class Qwen3TITOTokenizer(TITOTokenizer): + """Qwen3: the model stops at `<|im_end|>` without the trailing `\\n` the template + emits, so `merge_tokens` re-inserts it so the stored prefix stays canonical.""" + + reasoning_parser = "qwen3" + tool_call_parser = "qwen25" + _default_assistant_start_str = "<|im_start|>assistant" + + def __init__( + self, + tokenizer: Any, + chat_template_kwargs: dict[str, Any] | None = None, + assistant_start_str: str | None = None, + allowed_append_roles: list[str] | None = None, + ) -> None: + super().__init__( + tokenizer, + chat_template_kwargs, + assistant_start_str or self._default_assistant_start_str, + allowed_append_roles=allowed_append_roles, + ) + nl_ids = tokenizer.encode("\n", add_special_tokens=False) + if len(nl_ids) != 1: + raise ValueError(f"expected a single newline token, got {nl_ids}") + self._newline_id: int = nl_ids[0] + self._im_end_id: int = tokenizer.convert_tokens_to_ids("<|im_end|>") + self.trailing_token_ids = frozenset({self._newline_id}) + + def merge_tokens( + self, + old_messages: list[dict[str, Any]], + new_messages: list[dict[str, Any]], + pretokenized_token_ids: list[int], + tools: list[dict[str, Any]] | None = None, + ) -> list[int]: + incremental = self.tokenize_additional_non_assistant(old_messages, new_messages, tools) + prefix = list(pretokenized_token_ids) + if prefix and prefix[-1] == self._im_end_id: + prefix.append(self._newline_id) + return prefix + incremental + + +_QWEN3_FIXED = "qwen3_fixed.jinja" + + +def get_tito_tokenizer( + tokenizer: Any, + tokenizer_type: str = "qwen3", + *, + allowed_append_roles: tuple[str, ...] = ("tool",), +) -> TITOTokenizer: + """Build a TITO tokenizer. `default` uses the tokenizer's own chat template + (model-agnostic); `qwen3` loads the bundled fixed template (and disables thinking + clearing when `user` appends are allowed, so earlier turns keep their reasoning).""" + if tokenizer is None: + raise ValueError("tokenizer must not be None") + roles = frozenset(allowed_append_roles) + invalid = roles - _VALID_ROLES + if invalid: + raise ValueError(f"unknown roles in allowed_append_roles: {sorted(invalid)}; valid: {sorted(_VALID_ROLES)}") + + if tokenizer_type == "default": + return TITOTokenizer(tokenizer, allowed_append_roles=list(allowed_append_roles)) + if tokenizer_type == "qwen3": + kw: dict[str, Any] = {"chat_template": (TEMPLATE_DIR / _QWEN3_FIXED).read_text()} + if "user" in roles: + kw["clear_thinking"] = False + return Qwen3TITOTokenizer(tokenizer, chat_template_kwargs=kw, allowed_append_roles=list(allowed_append_roles)) + raise ValueError(f"unsupported tokenizer_type {tokenizer_type!r}; supported: 'qwen3', 'default'") diff --git a/plugins/tito/agentix/tito/engine/processing.py b/plugins/tito/agentix/tito/engine/processing.py new file mode 100644 index 0000000..578c28b --- /dev/null +++ b/plugins/tito/agentix/tito/engine/processing.py @@ -0,0 +1,16 @@ +"""Tokenizer loading. Minimal: load an HF tokenizer (tokenizer-only is fine — no +torch needed) and optionally override its chat template from a file.""" + +from __future__ import annotations + +from typing import Any + + +def load_tokenizer(name_or_path: str, chat_template_path: str | None = None, *, trust_remote_code: bool = False) -> Any: + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(name_or_path, trust_remote_code=trust_remote_code) + if chat_template_path: + with open(chat_template_path) as f: + tokenizer.chat_template = f.read() + return tokenizer diff --git a/plugins/tito/agentix/tito/engine/render.py b/plugins/tito/agentix/tito/engine/render.py new file mode 100644 index 0000000..ee69863 --- /dev/null +++ b/plugins/tito/agentix/tito/engine/render.py @@ -0,0 +1,94 @@ +"""Chat-template rendering backend. + +`apply_chat_template` renders messages through an HF tokenizer's chat template +(optionally an explicit `chat_template=` string for the fixed template), the same +code path SGLang uses. Tool definitions are canonicalized to the OpenAI +`{type:"function", function:{...}}` shape. No sglang dependency — the one pydantic +`Tool` type the canonicalization needs is defined locally. +""" + +from __future__ import annotations + +import copy +import json +from typing import Any, Literal + +from jinja2 import TemplateError +from pydantic import BaseModel, TypeAdapter + + +class _Function(BaseModel): + name: str + description: str | None = None + parameters: dict[str, Any] | None = None + + +class Tool(BaseModel): + type: str = "function" + function: _Function + + +def normalize_tool_arguments(messages: list[dict], format: Literal["dict", "json"]) -> list[dict]: + """Deep-copy *messages*, set assistant `content: None` -> "", and coerce tool_call + `arguments` to the form the renderer needs: "dict" (JSON string -> dict, for + HF-Jinja templates) or "json" (dict -> JSON string). Never mutates the input.""" + normalized = copy.deepcopy(messages) + for msg in normalized: + if msg.get("role") == "assistant": + if msg.get("content") is None: + msg["content"] = "" + if isinstance(msg.get("tool_calls"), list): + for item in msg["tool_calls"]: + func = item.get("function") + if not func: + continue + args = func.get("arguments") + if format == "dict" and isinstance(args, str): + func["arguments"] = json.loads(args) + elif format == "json" and isinstance(args, dict): + func["arguments"] = json.dumps(args, ensure_ascii=False) + return normalized + + +def extract_tool_dicts(tools: list[dict] | None) -> list[dict] | None: + """Canonicalize tools to full `{type:"function", function:{...}}` dumps.""" + if not tools: + return None + wrapped = [t if isinstance(t, dict) and "function" in t else {"type": "function", "function": t} for t in tools] + validated = TypeAdapter(list[Tool]).validate_python(wrapped) + return [tool.model_dump() for tool in validated] + + +def apply_chat_template( + messages: list[dict], + *, + tokenizer: Any, + tools: list[dict] | None = None, + add_generation_prompt: bool = True, + tokenize: bool = False, + **kwargs: Any, +) -> str | list[int]: + """Render via the HF tokenizer in SGLang style (`return_dict=False`, so the result + is `str` when tokenize=False or `list[int]` when tokenize=True). `chat_template=` + and other extras pass through `**kwargs`. Falls back to the bare function schema if + the template can't take the wrapped tool dicts.""" + messages = normalize_tool_arguments(messages, "dict") + tool_defs = extract_tool_dicts(tools) + render_kwargs = dict(add_generation_prompt=add_generation_prompt, **kwargs) + try: + return tokenizer.apply_chat_template( + messages, tokenize=tokenize, tools=tool_defs, return_dict=False, **render_kwargs + ) + except TemplateError as e: + if tool_defs is not None: + try: + return tokenizer.apply_chat_template( + messages, + tokenize=tokenize, + tools=[t["function"] if "function" in t else t for t in tool_defs], + return_dict=False, + **render_kwargs, + ) + except TemplateError as te: + raise ValueError(f"Chat template rendering failed (tool format fallback): {te}") from te + raise ValueError(f"Chat template rendering failed: {e}") from e diff --git a/plugins/tito/agentix/tito/engine/session_app.py b/plugins/tito/agentix/tito/engine/session_app.py new file mode 100644 index 0000000..d76709e --- /dev/null +++ b/plugins/tito/agentix/tito/engine/session_app.py @@ -0,0 +1,231 @@ +"""FastAPI session routes for the TITO gateway. + +The gateway keeps a token-aligned trajectory per session and proxies chat +completions to an OpenAI-compatible backend (sglang). The chat-completions flow: +prepare pretokenized input_ids (lock held briefly) -> force logprobs/meta_info -> +proxy to the backend (no lock) -> validate -> append the trajectory checkpoint +(lock held briefly). The proxy is NOT held under the lock so a slow generation +doesn't block DELETE/other ops. + +`build_session_app` is backend-agnostic: pass any object exposing +``do_proxy(request, path, body=None) -> dict`` and ``build_proxy_response(result)``. +""" + +from __future__ import annotations + +import json +import logging +import time +from typing import Any, Protocol + +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse +from starlette.responses import Response + +from .errors import ( + MessageValidationError, + SessionError, + SessionNotFoundError, + TokenizationError, + UpstreamResponseError, +) +from .pretokenize import get_tito_tokenizer +from .processing import load_tokenizer +from .trajectory import GetSessionResponse, SessionRecord, SessionRegistry + +logger = logging.getLogger(__name__) + + +class Backend(Protocol): + async def do_proxy(self, request: Request, path: str, body: bytes | None = None) -> dict: ... + def build_proxy_response(self, result: dict) -> Response: ... + + +def build_registry(args: Any) -> SessionRegistry | None: + """Construct a SessionRegistry from gateway args, or None if no hf_checkpoint.""" + hf_checkpoint = getattr(args, "hf_checkpoint", None) + if not hf_checkpoint: + logger.info("[session] no hf_checkpoint set — session routes disabled") + return None + tokenizer = load_tokenizer( + hf_checkpoint, + chat_template_path=getattr(args, "chat_template_path", None), + # Executes Python shipped inside the checkpoint repo — explicit opt-in + # (config/--trust-remote-code), never a hardcoded default. + trust_remote_code=bool(getattr(args, "trust_remote_code", False)), + ) + roles = getattr(args, "tito_allowed_append_roles", None) or ("tool",) + tito_tokenizer = get_tito_tokenizer( + tokenizer, + tokenizer_type=getattr(args, "tito_model", "default"), + allowed_append_roles=tuple(roles), + ) + return SessionRegistry(args, tokenizer, tito_tokenizer=tito_tokenizer) + + +def setup_session_routes(app: FastAPI, backend: Backend, args: Any) -> None: + registry = build_registry(args) + if registry is None: + return + # Exposed for operational introspection (session counts, tests). + app.state.tito_registry = registry + + instance_id = getattr(args, "session_server_instance_id", None) + + @app.exception_handler(SessionError) + async def _session_error_handler(request: Request, exc: SessionError) -> JSONResponse: + return JSONResponse(status_code=exc.status_code, content={"error": str(exc)}) + + @app.get("/health") + async def health() -> dict[str, Any]: + body: dict[str, Any] = {"status": "ok"} + if instance_id is not None: + body["session_server_instance_id"] = instance_id + return body + + @app.post("/sessions") + async def create_session() -> dict[str, str]: + return {"session_id": registry.create_session()} + + @app.get("/sessions/{session_id}") + async def get_session(session_id: str) -> GetSessionResponse: + session = registry.get_session(session_id) + metadata: dict[str, Any] = {} + try: + mismatch = registry.compute_session_mismatch(session) + except TokenizationError: + logger.exception("failed to compute tito_session_mismatch for %s", session_id) + mismatch = None + if mismatch is not None: + metadata["tito_session_mismatch"] = mismatch + metadata["accumulated_token_ids"] = session.token_ids + metadata["max_trim_tokens"] = registry.tito_tokenizer.max_trim_tokens + return GetSessionResponse(session_id=session_id, records=session.records, metadata=metadata) + + @app.delete("/sessions/{session_id}") + async def delete_session(session_id: str) -> Response: + session = registry.get_session(session_id) + if session.closing: + raise SessionNotFoundError(f"session not found: session_id={session_id}") + # `closing` is only ever set INSIDE the lock, immediately before the + # removal: a DELETE cancelled while waiting for the lock (client + # timeout/disconnect) must leave the session deletable, not wedged + # behind a stuck flag with the trajectory leaked in the registry. + async with session.lock: + if session.closing: + raise SessionNotFoundError(f"session not found: session_id={session_id}") + session.closing = True + registry.remove_session(session_id) + return Response(status_code=204) + + @app.post("/sessions/{session_id}/v1/chat/completions") + async def chat_completions(request: Request, session_id: str) -> Response: + session = registry.get_session(session_id) + if session.closing: + raise SessionNotFoundError(f"session not found: session_id={session_id}") + + # Read + parse the body BEFORE taking the lock: the read lasts as long + # as the client's upload — under the lock, one dribbling client would + # wedge DELETE and every other operation on the session. + raw = await request.body() + try: + request_body = json.loads(raw) if raw else {} + except ValueError as e: + raise MessageValidationError(f"request body is not valid JSON: {e}") from e + if not isinstance(request_body, dict): + raise MessageValidationError("request body must be a JSON object") + + # Phase 1: prepare pretokenized input_ids (lock held briefly). + async with session.lock: + if session.closing: + raise SessionNotFoundError(f"session not found: session_id={session_id}") + # Hardcoded so an agent override can't break token accumulation: + request_body["logprobs"] = True # -> meta_info.output_token_logprobs + request_body["return_meta_info"] = True # -> choice.meta_info + request_body["no_stop_trim"] = False # stop-token text trimmed from content + # The TITO flow needs the complete JSON completion (logprobs + + # meta_info); an SSE stream would be unparseable below. Force + # non-streaming — a stream:true agent gets the full JSON body back. + request_body["stream"] = False + request_body.pop("stream_options", None) + request_messages = request_body.get("messages", []) + prompt_token_ids = session.prepare_pretokenized( + request_messages, tools=request_body.get("tools"), tito_tokenizer=registry.tito_tokenizer + ) + request_body["input_ids"] = prompt_token_ids + body = json.dumps(request_body).encode() + # `version` advances on BOTH update and rollback; `num_assistant` + # would be ABA-prone (a concurrent rollback+update restores it). + expected_version = session.version + + # Phase 2: proxy to the backend (NO lock). + result = await backend.do_proxy(request, "v1/chat/completions", body=body) + if result["status_code"] != 200: + return backend.build_proxy_response(result) + + # Structural failures in a 200 body are the backend's fault — surface + # every shape violation as a clean 502, never an unhandled 500. + try: + response = json.loads(result["response_body"]) + except ValueError as e: + raise UpstreamResponseError(f"backend returned an unparseable 200 body: {e}") from e + if not isinstance(response, dict): + raise UpstreamResponseError("backend 200 body is not a JSON object") + choices = response.get("choices") + if not isinstance(choices, list) or not choices or not isinstance(choices[0], dict): + raise UpstreamResponseError("backend 200 body has no choices") + choice = choices[0] + meta_info = choice.get("meta_info") + if not isinstance(meta_info, dict) or "output_token_logprobs" not in meta_info: + raise UpstreamResponseError("meta_info.output_token_logprobs missing (needs logprobs=True)") + assistant_message = choice.get("message") + if not isinstance(assistant_message, dict): + raise UpstreamResponseError("assistant message missing") + if assistant_message.get("content") is None and not assistant_message.get("tool_calls"): + # Tool-call-only turns routinely carry content:null (the parser + # consumed all generated text) — only a turn with NEITHER content + # NOR tool_calls is malformed. + raise UpstreamResponseError("assistant message has neither content nor tool_calls") + output_token_logprobs = meta_info["output_token_logprobs"] + completion_tokens = meta_info.get("completion_tokens") + if not isinstance(output_token_logprobs, list) or not isinstance(completion_tokens, int): + raise UpstreamResponseError("meta_info output_token_logprobs/completion_tokens malformed") + if len(output_token_logprobs) != completion_tokens: + raise UpstreamResponseError( + f"len(output_token_logprobs)={len(output_token_logprobs)} != completion_tokens={completion_tokens}" + ) + try: + completion_token_ids = [t[1] for t in output_token_logprobs] + except (TypeError, IndexError, KeyError) as e: + raise UpstreamResponseError(f"malformed output_token_logprobs entry: {e}") from e + + # Phase 3: append the trajectory checkpoint (lock held briefly). + async with session.lock: + if session.closing: + return backend.build_proxy_response(result) + if session.version != expected_version: + logger.warning("session %s changed during proxy; skipping state update", session_id) + return backend.build_proxy_response(result) + session.update_pretokenized_state( + request_messages, + assistant_message, + prompt_token_ids=prompt_token_ids, + completion_token_ids=completion_token_ids, + max_trim_tokens=registry.tito_tokenizer.max_trim_tokens, + ) + session.append_record( + SessionRecord( + timestamp=time.time(), + method=request.method, + path="/v1/chat/completions", + status_code=result["status_code"], + request=request_body, + response=response, + ) + ) + return backend.build_proxy_response(result) + + @app.api_route("/sessions/{session_id}/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"]) + async def session_proxy(request: Request, session_id: str, path: str) -> Response: + result = await backend.do_proxy(request, path) + return backend.build_proxy_response(result) diff --git a/plugins/tito/agentix/tito/engine/templates/qwen3_fixed.jinja b/plugins/tito/agentix/tito/engine/templates/qwen3_fixed.jinja new file mode 100644 index 0000000..88ca535 --- /dev/null +++ b/plugins/tito/agentix/tito/engine/templates/qwen3_fixed.jinja @@ -0,0 +1,85 @@ +{%- if tools %} + {{- '<|im_start|>system\n' }} + {%- if messages[0].role == 'system' %} + {{- messages[0].content + '\n\n' }} + {%- endif %} + {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within XML tags:\n" }} + {%- for tool in tools %} + {{- "\n" }} + {{- tool | tojson }} + {%- endfor %} + {{- "\n\n\nFor each function call, return a json object with function name and arguments within XML tags:\n\n{\"name\": , \"arguments\": }\n<|im_end|>\n" }} +{%- else %} + {%- if messages[0].role == 'system' %} + {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }} + {%- endif %} +{%- endif %} +{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %} +{%- for message in messages[::-1] %} + {%- set index = (messages|length - 1) - loop.index0 %} + {%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('') and message.content.endswith('')) %} + {%- set ns.multi_step_tool = false %} + {%- set ns.last_query_index = index %} + {%- endif %} +{%- endfor %} +{%- for message in messages %} + {%- if message.content is string %} + {%- set content = message.content %} + {%- else %} + {%- set content = '' %} + {%- endif %} + {%- if (message.role == "user") or (message.role == "system" and not loop.first) %} + {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }} + {%- elif message.role == "assistant" %} + {%- set reasoning_content = '' %} + {%- if message.reasoning_content is string %} + {%- set reasoning_content = message.reasoning_content %} + {%- else %} + {%- if '' in content %} + {%- set reasoning_content = content.split('')[0].rstrip('\n').split('')[-1].lstrip('\n') %} + {%- set content = content.split('')[-1].lstrip('\n') %} + {%- endif %} + {%- endif %} + {%- if (not (clear_thinking | default(true))) or loop.index0 > ns.last_query_index %} + {{- '<|im_start|>' + message.role + '\n\n' + reasoning_content.strip('\n') + '\n\n\n' + content.lstrip('\n') }} + {%- else %} + {{- '<|im_start|>' + message.role + '\n' + content }} + {%- endif %} + {%- if message.tool_calls %} + {%- for tool_call in message.tool_calls %} + {%- if (loop.first and content) or (not loop.first) %} + {{- '\n' }} + {%- endif %} + {%- if tool_call.function %} + {%- set tool_call = tool_call.function %} + {%- endif %} + {{- '\n{"name": "' }} + {{- tool_call.name }} + {{- '", "arguments": ' }} + {%- if tool_call.arguments is string %} + {{- tool_call.arguments }} + {%- else %} + {{- tool_call.arguments | tojson }} + {%- endif %} + {{- '}\n' }} + {%- endfor %} + {%- endif %} + {{- '<|im_end|>\n' }} + {%- elif message.role == "tool" %} + {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %} + {{- '<|im_start|>user' }} + {%- endif %} + {{- '\n\n' }} + {{- content }} + {{- '\n' }} + {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %} + {{- '<|im_end|>\n' }} + {%- endif %} + {%- endif %} +{%- endfor %} +{%- if add_generation_prompt %} + {{- '<|im_start|>assistant\n' }} + {%- if enable_thinking is defined and enable_thinking is false %} + {{- '\n\n\n\n' }} + {%- endif %} +{%- endif %} diff --git a/plugins/tito/agentix/tito/engine/trajectory.py b/plugins/tito/agentix/tito/engine/trajectory.py new file mode 100644 index 0000000..1e1ac43 --- /dev/null +++ b/plugins/tito/agentix/tito/engine/trajectory.py @@ -0,0 +1,275 @@ +"""Linear trajectory state machine + session registry. + +`LinearTrajectory` holds one session's message history and accumulated token-ID +checkpoints, and is the heart of incremental pretokenization: on each turn it +validates that the request extends the stored history (rolling back at most one +assistant step on agent retries) and reuses the stored token prefix. `SessionRegistry` +maps session IDs to trajectories and computes the from-scratch-vs-accumulated +mismatch report. Mutating methods must be called under `LinearTrajectory.lock`. +""" + +from __future__ import annotations + +import asyncio +import logging +import uuid +from dataclasses import dataclass, field +from typing import Any + +from pydantic import BaseModel, Field + +from .compare import TokenSeqComparator +from .errors import MessageValidationError, SessionNotFoundError, TokenizationError +from .messages import assert_messages_append_only_with_allowed_role, message_matches +from .pretokenize import TITOTokenizer + +logger = logging.getLogger(__name__) + +# Only single-step rollback is supported (an agent retrying one tool call). +MAX_ASSISTANT_ROLLBACK_STEPS = 1 + + +class SessionRecord(BaseModel): + timestamp: float + method: str + path: str + request: dict + response: dict + status_code: int + + +class GetSessionResponse(BaseModel): + session_id: str + records: list[SessionRecord] + metadata: dict = Field(default_factory=dict) + + +@dataclass(frozen=True, slots=True) +class _RollbackPlan: + """A computed-but-not-applied rollback: message truncation point, the + checkpoint to land on, and how many trailing checkpoints to discard.""" + + msg_end: int + checkpoint_index: int + discard_count: int + + +@dataclass +class LinearTrajectory: + """Message history + accumulated token-ID checkpoints for one session.""" + + lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False, compare=False) + closing: bool = field(default=False, repr=False, compare=False) + messages: list[dict[str, Any]] = field(default_factory=list) + records: list[SessionRecord] = field(default_factory=list) + trajectory_token_ids: list[list[int]] = field(default_factory=list) + num_assistant: int = 0 + # Monotonic change counter, bumped by BOTH update and rollback. Unlike + # `num_assistant` (which rollback decrements and update restores — an ABA + # hazard), an equal `version` proves the session was untouched in between. + version: int = 0 + + @property + def token_ids(self) -> list[int]: + """The latest assistant checkpoint's token IDs.""" + return self.trajectory_token_ids[-1] if self.trajectory_token_ids else [] + + def append_record(self, record: SessionRecord) -> None: + self.records.append(record) + + def prepare_pretokenized( + self, + request_messages: list[dict[str, Any]], + tools: list[dict[str, Any]] | None = None, + *, + tito_tokenizer: TITOTokenizer, + ) -> list[int]: + """Build the full prompt token IDs for *request_messages*. First turn renders + from scratch; later turns reuse the stored token prefix (rolling back at most + one assistant step on a retry). Must be called under ``self.lock``.""" + if not self.messages: + return tito_tokenizer.render_messages( + request_messages, tools=tools, add_generation_prompt=True, tokenize=True + ) + + # Plan first, mutate last: a request that fails validation must be a + # pure 4xx with NO committed rollback side effects — otherwise a + # rejected request silently truncates the trajectory and bricks the + # original branch. + plan = self._plan_rollback(request_messages) + base_messages = self.messages if plan is None else self.messages[: plan.msg_end] + try: + assert_messages_append_only_with_allowed_role( + base_messages, request_messages, tito_tokenizer.allowed_append_roles + ) + except ValueError as e: + raise MessageValidationError(f"{e}; to allow more roles use --tito-allowed-append-roles") from e + if plan is not None: + self._apply_rollback(plan) + + if not self.token_ids: + # The token prefix is unavailable — either a retry chain walked + # past the trimmed checkpoint window, or a prior turn never landed + # its update. Re-render from scratch: incremental == from-scratch + # is the engine invariant, so this is lossless (never merge onto + # an empty prefix, which would drop the whole stored history). + return tito_tokenizer.render_messages( + request_messages, tools=tools, add_generation_prompt=True, tokenize=True + ) + return tito_tokenizer.merge_tokens( + old_messages=self.messages, + new_messages=request_messages, + pretokenized_token_ids=self.token_ids, + tools=tools, + ) + + def update_pretokenized_state( + self, + request_messages: list[dict[str, Any]], + assistant_message: dict[str, Any], + prompt_token_ids: list[int], + completion_token_ids: list[int], + max_trim_tokens: int, + ) -> None: + """Append ``prompt+completion`` token IDs as a new checkpoint after a successful + response, validating the previously-stored IDs are a prefix (tolerating up to + ``max_trim_tokens`` trailing differences). Must be called under ``self.lock``.""" + all_token_ids = prompt_token_ids + completion_token_ids + + prev = self.token_ids + if prev: + check_len = len(prev) - max_trim_tokens + if check_len > 0 and all_token_ids[:check_len] != prev[:check_len]: + # strict=False: the new sequence may legitimately be SHORTER + # than the stored checkpoint (that length difference IS the + # mismatch) — the default below covers the exhausted tail. + first_mismatch = next( + ( + i + for i, (a, b) in enumerate(zip(all_token_ids[:check_len], prev[:check_len], strict=False)) + if a != b + ), + min(len(all_token_ids), check_len), + ) + raise TokenizationError( + f"pretokenized prefix mismatch: stored {len(prev)} tokens " + f"(checking first {check_len}, allowing {max_trim_tokens} trailing) are not a prefix of " + f"prompt_token_ids + completion_token_ids ({len(all_token_ids)} tokens), " + f"first mismatch at index {first_mismatch}, matched {first_mismatch}/{check_len} prefix tokens\n" + f"request_messages={request_messages}\nassistant_message={assistant_message}" + ) + + self.messages = list(request_messages) + [assistant_message] + self.trajectory_token_ids.append(all_token_ids) + self.num_assistant += 1 + self.version += 1 + # Single-step rollback (MAX_ASSISTANT_ROLLBACK_STEPS=1) can only ever + # reach the last two checkpoints; keeping every full prefix+completion + # list would be O(turns^2) dead memory per session. (A retry chain + # that outruns this window falls back to a from-scratch render in + # `prepare_pretokenized`.) + if len(self.trajectory_token_ids) > MAX_ASSISTANT_ROLLBACK_STEPS + 1: + del self.trajectory_token_ids[: -(MAX_ASSISTANT_ROLLBACK_STEPS + 1)] + + def _plan_rollback(self, request_messages: list[dict[str, Any]]) -> _RollbackPlan | None: + """If *request_messages* diverges from the stored history, compute the + rollback to the last assistant checkpoint within the matching prefix + (single-step only). Pure — mutates nothing; `None` means no divergence.""" + stored = self.messages + if not stored: + return None + + match_len = 0 + for i in range(min(len(request_messages), len(stored))): + if message_matches(stored[i], request_messages[i]): + match_len = i + 1 + else: + break + + if match_len >= len(stored): + return None + + rollback_msg_end = 0 + checkpoint_index = -1 + assistant_count = 0 + for i in range(match_len): + if stored[i].get("role") == "assistant": + rollback_msg_end = i + 1 + checkpoint_index = assistant_count + assistant_count += 1 + + if checkpoint_index < 0: + raise MessageValidationError( + f"rollback failed: no assistant message found in the first {match_len} matched messages " + f"(stored has {len(stored)} messages, request has {len(request_messages)} messages)" + ) + + discard_count = self.num_assistant - (checkpoint_index + 1) + if discard_count > MAX_ASSISTANT_ROLLBACK_STEPS: + raise MessageValidationError( + f"rollback failed: discard_count={discard_count} exceeds " + f"max_assistant_rollback_steps={MAX_ASSISTANT_ROLLBACK_STEPS} " + f"(stored has {len(stored)} messages, request has {len(request_messages)} messages)" + ) + return _RollbackPlan( + msg_end=rollback_msg_end, checkpoint_index=checkpoint_index, discard_count=discard_count + ) + + def _apply_rollback(self, plan: _RollbackPlan) -> None: + logger.info( + "Rolling back session: stored %d messages / %d assistants -> checkpoint %d (messages[:%d]), " + "discarding %d assistant(s)", + len(self.messages), self.num_assistant, plan.checkpoint_index, plan.msg_end, plan.discard_count, + ) + self.messages = self.messages[: plan.msg_end] + # trajectory_token_ids holds only the trailing checkpoints (older ones + # are trimmed as unreachable), so discard relative to the end — the + # absolute checkpoint_index no longer maps onto the list. A retry chain + # can legally outrun the trimmed window: clamp, possibly to empty (the + # caller then re-renders from scratch instead of merging). + if plan.discard_count and self.trajectory_token_ids: + del self.trajectory_token_ids[-min(plan.discard_count, len(self.trajectory_token_ids)):] + self.records = self.records[: plan.checkpoint_index + 1] + self.num_assistant = plan.checkpoint_index + 1 + self.version += 1 + + +class SessionRegistry: + """Session ID -> trajectory map + shared tokenizer/comparator. Pure CRUD plus the + read-only mismatch computation; never mutates trajectory state itself.""" + + def __init__(self, args: Any, tokenizer: Any, *, tito_tokenizer: TITOTokenizer) -> None: + self.sessions: dict[str, LinearTrajectory] = {} + self.args = args + self.tokenizer = tokenizer + self.tito_tokenizer = tito_tokenizer + self.comparator: TokenSeqComparator = tito_tokenizer.create_comparator() + + def create_session(self) -> str: + session_id = uuid.uuid4().hex + self.sessions[session_id] = LinearTrajectory() + return session_id + + def get_session(self, session_id: str) -> LinearTrajectory: + session = self.sessions.get(session_id) + if session is None: + raise SessionNotFoundError(f"session not found: session_id={session_id}") + return session + + def remove_session(self, session_id: str) -> None: + if self.sessions.pop(session_id, None) is None: + raise SessionNotFoundError(f"session not found: session_id={session_id}") + + def compute_session_mismatch(self, session: LinearTrajectory) -> list[dict] | None: + """Compare accumulated token IDs against a from-scratch render. Read-only.""" + if not session.token_ids: + return None + try: + tools = session.records[-1].request.get("tools") if session.records else None + expected_ids = self.tito_tokenizer.render_messages( + session.messages, tools=tools, add_generation_prompt=False, tokenize=True + ) + mismatches = self.comparator.compare_sequences(expected_ids, session.token_ids) + return [m.to_dict() for m in mismatches] + except Exception as e: + raise TokenizationError(f"failed to compute tito_session_mismatch: {e}") from e diff --git a/plugins/tito/agentix/tito/gateway.py b/plugins/tito/agentix/tito/gateway.py new file mode 100644 index 0000000..942bcb9 --- /dev/null +++ b/plugins/tito/agentix/tito/gateway.py @@ -0,0 +1,61 @@ +"""Python wrapper API for launching TITO Gateway beside a backend server.""" + +from __future__ import annotations + +from dataclasses import replace + +from .config import TITOGatewayConfig +from .discovery import discover_backend_url, normalize_backend_url +from .pool import BackendPool +from .server import SessionServer + + +class TITOGateway: + """Small wrapper that resolves backend(s) and owns a session server app. + + Routes inference across a :class:`BackendPool` — a single backend (resolved + by discovery) by default, or several when ``config.backend_urls`` is set. + """ + + def __init__(self, config: TITOGatewayConfig): + if config.backend_urls: + urls = [normalize_backend_url(u) for u in config.backend_urls] + self.config = replace(config, backend_url=urls[0]) + else: + backend_url = discover_backend_url( + config.backend_url, + probe_candidates=config.backend_probe_candidates, + probe_timeout=config.backend_probe_timeout, + ) + self.config = replace(config, backend_url=backend_url) + urls = [backend_url] + self.pool = BackendPool(urls, policy=config.routing_policy) + self.server = SessionServer(self.config.as_session_args(), self.pool) + self._register_health_alias() + + @classmethod + def from_server(cls, *, hf_checkpoint: str, backend_url: str | None = None, **kwargs) -> TITOGateway: + return cls(TITOGatewayConfig(hf_checkpoint=hf_checkpoint, backend_url=backend_url, **kwargs)) + + def _register_health_alias(self) -> None: + # abridge's Sidecar probes `/healthz` by default; the engine session + # routes only expose `/health`. Add a thin alias so a default Sidecar + # wiring works without overriding `health_path`. + async def healthz() -> dict[str, str]: + return {"status": "ok"} + + self.app.add_api_route("/healthz", healthz, methods=["GET"]) + + @property + def app(self): + return self.server.app + + def run(self) -> None: + import uvicorn + + uvicorn.run( + self.app, + host=self.config.session_server_ip, + port=self.config.session_server_port, + log_level="info", + ) diff --git a/plugins/tito/agentix/tito/pool.py b/plugins/tito/agentix/tito/pool.py new file mode 100644 index 0000000..93fa480 --- /dev/null +++ b/plugins/tito/agentix/tito/pool.py @@ -0,0 +1,119 @@ +"""Backend pool — route OpenAI-compatible requests across N base URLs. + +The TITO Gateway accepts one *or more* OpenAI-compatible backend URLs +(sglang replicas — the token-recording chat flow needs sglang's ``meta_info`` +extension; the raw proxy routes work with any OpenAI-compatible server) and +forwards each request to one of them. This is the routing layer, independent +of TITO tokenization, so it is unit-tested on its own with no model in the +loop. + +Policy: + - ``sticky`` (default): each ``session_id`` is pinned to one backend, + chosen round-robin among healthy backends on first sight and then + remembered. A multi-turn rollout reuses one replica's prefix KV-cache, + which is the right default for TITO. (TITO sends explicit ``input_ids``, + so any replica *can* serve any turn — stickiness is a cache-locality + optimization, not a correctness requirement.) + - ``round_robin``: spread every request across healthy backends. + +Backends reported down via ``report_down`` are skipped; a sticky session +whose backend goes down is reassigned on its next pick. Recovery is +two-pronged: the proxy layer calls ``report_up`` whenever a backend answers +a request, and a down mark expires after ``down_cooldown`` seconds anyway +(half-open: the backend becomes eligible again; if it is still broken the +next failed request re-marks it). Sticky pins are LRU-bounded at +``max_pins`` — the recommended rollout flow never DELETEs its session (the +trajectory must survive for harvest), so an unbounded pin map would grow +forever. Losing an old pin only costs prefix-cache locality, never +correctness. +""" + +from __future__ import annotations + +import threading +import time +from collections import OrderedDict +from collections.abc import Sequence + +_POLICIES = ("sticky", "round_robin") + + +class BackendPool: + def __init__( + self, + backends: Sequence[str], + *, + policy: str = "sticky", + down_cooldown: float = 30.0, + max_pins: int = 10_000, + ) -> None: + urls = [b.rstrip("/") for b in backends if b] + if not urls: + raise ValueError("BackendPool requires at least one backend url") + if policy not in _POLICIES: + raise ValueError(f"policy must be one of {_POLICIES}; got {policy!r}") + if max_pins < 1: + raise ValueError(f"max_pins must be >= 1; got {max_pins}") + self._backends = urls + self._policy = policy + self._rr = 0 + self._assigned: OrderedDict[str, str] = OrderedDict() + self._down: dict[str, float] = {} # url -> monotonic time marked down + self._down_cooldown = down_cooldown + self._max_pins = max_pins + self._lock = threading.Lock() + + @property + def backends(self) -> tuple[str, ...]: + return tuple(self._backends) + + def _is_down(self, backend: str, now: float) -> bool: + marked = self._down.get(backend) + return marked is not None and (now - marked) < self._down_cooldown + + def _healthy(self, now: float) -> list[str]: + healthy = [b for b in self._backends if not self._is_down(b, now)] + # All down → fall back to the full set rather than fail the request; + # the forward attempt surfaces the real error. + return healthy or list(self._backends) + + def _next_round_robin(self, healthy: list[str]) -> str: + chosen = healthy[self._rr % len(healthy)] + self._rr += 1 + return chosen + + def pick(self, session_id: str | None = None) -> str: + """Choose a backend for a request. With the sticky policy and a + `session_id`, return that session's pinned backend (assigning one + the first time, or reassigning if the pinned one is down).""" + with self._lock: + now = time.monotonic() + healthy = self._healthy(now) + if self._policy == "sticky" and session_id is not None: + current = self._assigned.get(session_id) + if current is not None and not self._is_down(current, now): + self._assigned.move_to_end(session_id) + return current + chosen = self._next_round_robin(healthy) + self._assigned[session_id] = chosen + self._assigned.move_to_end(session_id) + while len(self._assigned) > self._max_pins: + self._assigned.popitem(last=False) + return chosen + return self._next_round_robin(healthy) + + def report_down(self, backend: str) -> None: + with self._lock: + self._down[backend.rstrip("/")] = time.monotonic() + + def report_up(self, backend: str) -> None: + with self._lock: + self._down.pop(backend.rstrip("/"), None) + + def forget(self, session_id: str) -> None: + """Drop a session's sticky assignment (call when the rollout ends).""" + with self._lock: + self._assigned.pop(session_id, None) + + +__all__ = ["BackendPool"] diff --git a/plugins/tito/agentix/tito/server.py b/plugins/tito/agentix/tito/server.py new file mode 100644 index 0000000..cceb63a --- /dev/null +++ b/plugins/tito/agentix/tito/server.py @@ -0,0 +1,118 @@ +"""Session server — a FastAPI app over the native TITO engine, routing proxied +inference across a multi-backend pool. + +The engine's `session_app` owns the routes (sessions + the token-aligned chat +flow); this module supplies the *backend*: a pooled httpx proxy that picks a +replica per request (sticky by ``session_id`` for prefix-cache locality), reports +a replica down on a transport error, and forgets a session's pin on delete. +""" + +from __future__ import annotations + +import json +import logging +from typing import Any + +import httpx +from fastapi import FastAPI, Request +from starlette.responses import Response + +from .engine.session_app import setup_session_routes +from .pool import BackendPool + +logger = logging.getLogger(__name__) + +_HOP_BY_HOP = ("content-length", "transfer-encoding", "host") +_RESP_STRIP = ("content-length", "transfer-encoding", "content-encoding") + + +def _session_id_from_path(path: str) -> str | None: + """Extract ``{session_id}`` from ``/sessions/{session_id}[/...]``.""" + parts = path.strip("/").split("/") + if len(parts) >= 2 and parts[0] == "sessions": + return parts[1] + return None + + +class _PooledBackend: + """Backend for the session routes: proxy each request to a pool-picked replica.""" + + def __init__(self, args: Any, pool: BackendPool) -> None: + self._pool = pool + timeout = getattr(args, "router_timeout", 600.0) + self.client = httpx.AsyncClient( + limits=httpx.Limits(max_connections=1024), timeout=httpx.Timeout(timeout) + ) + + async def do_proxy(self, request: Request, path: str, body: bytes | None = None) -> dict: + session_id = _session_id_from_path(request.url.path) + backend_url = self._pool.pick(session_id) + url = f"{backend_url}/{path}" + if request.url.query: + url = f"{url}?{request.url.query}" + if body is None: + body = await request.body() + headers = {k: v for k, v in request.headers.items() if k.lower() not in _HOP_BY_HOP} + try: + response = await self.client.request(request.method, url, content=body, headers=headers) + except httpx.TransportError as exc: + self._pool.report_down(backend_url) + logger.warning("pooled proxy transport error %s -> %s: %s", path, backend_url, exc) + error_body = json.dumps({"error": f"backend transport error: {type(exc).__name__}: {exc}"}).encode() + return { + "request_body": body, + "response_body": error_body, + "status_code": 502, + "headers": {"content-type": "application/json"}, + } + # Any HTTP response means the transport is alive — clear a stale down + # mark (picked here via the all-down fallback or a cooldown half-open). + self._pool.report_up(backend_url) + content = await response.aread() + return { + "request_body": body, + "response_body": content, + "status_code": response.status_code, + "headers": dict(response.headers), + } + + def build_proxy_response(self, result: dict) -> Response: + # Pass the upstream body through VERBATIM. Re-encoding via json would + # perturb key order/whitespace for byte-sensitive consumers and 500 on + # JSON that Python parses but a strict serializer rejects (NaN/Infinity). + content = result["response_body"] + headers = {k: v for k, v in result["headers"].items() if k.lower() not in _RESP_STRIP} + return Response( + content=content, + status_code=result["status_code"], + headers=headers, + media_type=headers.get("content-type", ""), + ) + + async def aclose(self) -> None: + await self.client.aclose() + + +class SessionServer: + """FastAPI session server backed by the native TITO engine + a BackendPool.""" + + def __init__(self, args: Any, pool: BackendPool) -> None: + self.args = args + self.pool = pool + self.backend_url = pool.backends[0] + self.app = FastAPI() + self._backend = _PooledBackend(args, pool) + self.app.router.on_shutdown.append(self._backend.aclose) + setup_session_routes(self.app, self._backend, args) + self.app.middleware("http")(self._forget_on_delete) + + async def _forget_on_delete(self, request: Request, call_next: Any) -> Response: + response = await call_next(request) + if request.method == "DELETE" and response.status_code < 300: + # Only the session resource itself (`DELETE /sessions/{id}`) ends + # the rollout — a DELETE proxied to a backend sub-resource under + # the session prefix must not drop a live session's sticky pin. + parts = request.url.path.strip("/").split("/") + if len(parts) == 2 and parts[0] == "sessions": + self.pool.forget(parts[1]) + return response diff --git a/plugins/tito/agentix/tito/tokenizer.py b/plugins/tito/agentix/tito/tokenizer.py new file mode 100644 index 0000000..bf3fa11 --- /dev/null +++ b/plugins/tito/agentix/tito/tokenizer.py @@ -0,0 +1,29 @@ +"""Public tokenizer entrypoints — thin re-export of the native TITO engine.""" + +from __future__ import annotations + +from enum import StrEnum +from typing import Any + +from .engine.pretokenize import get_tito_tokenizer as _engine_get_tito_tokenizer + + +class TITOTokenizerType(StrEnum): + """Tokenizer families the native engine supports. Other models are a small + subclass + a fixed chat template — see agentix.tito.engine.pretokenize.""" + + DEFAULT = "default" + QWEN3 = "qwen3" + + +def get_tito_tokenizer( + tokenizer: Any, + tokenizer_type: TITOTokenizerType | str = TITOTokenizerType.DEFAULT, + *, + allowed_append_roles: tuple[str, ...] | list[str] | None = None, + **_ignored: Any, +) -> Any: + """Build a TITO tokenizer for *tokenizer* (`"qwen3"` or `"default"`).""" + t = tokenizer_type.value if isinstance(tokenizer_type, TITOTokenizerType) else str(tokenizer_type) + roles = tuple(allowed_append_roles) if allowed_append_roles else ("tool",) + return _engine_get_tito_tokenizer(tokenizer, t, allowed_append_roles=roles) diff --git a/plugins/tito/pyproject.toml b/plugins/tito/pyproject.toml new file mode 100644 index 0000000..374e95d --- /dev/null +++ b/plugins/tito/pyproject.toml @@ -0,0 +1,69 @@ +[build-system] +requires = ["uv_build>=0.7,<0.9"] +build-backend = "uv_build" + +[project] +name = "agentix-tito" +version = "0.1.0" +description = "Agentix TITO plugin — token-in-token-out session-recording gateway." +readme = "README.md" +requires-python = ">=3.11" +license = "Apache-2.0" +authors = [ + { name = "Agentix maintainers" }, +] +keywords = ["agentix", "tito", "agentic", "chat-template", "gateway", "rollout"] +classifiers = [ + "Development Status :: 3 - Alpha", + "Environment :: Console", + "Framework :: FastAPI", + "Intended Audience :: Developers", + "License :: OSI Approved :: Apache Software License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Topic :: Internet :: WWW/HTTP :: HTTP Servers", + "Topic :: Scientific/Engineering :: Artificial Intelligence", +] +# The gateway tokenizes prompts itself (the native TITO engine — see +# agentix.tito.engine), so transformers + tokenizers + jinja2 are real runtime +# deps. They are isolated to THIS plugin — agentix core / abridge never pull +# them. No sglang dependency: the engine defines the one pydantic `Tool` type it +# needs itself. +dependencies = [ + # As an `agentix.tito` plugin it lives in the agentix namespace, so importing it + # runs agentix core's __init__ — hence the agentixx dep (pure plumbing: socketio, + # msgpack, fastapi; zero ML/training-framework code). + "agentixx", + "fastapi>=0.110", + "httpx>=0.27", + "pydantic>=2", + "uvicorn>=0.29", + "transformers>=4.44", + "tokenizers>=0.19", + "jinja2>=3.1", + "huggingface-hub>=0.23", +] + +[project.optional-dependencies] +test = ["pytest>=8", "pytest-asyncio>=0.23"] + +[project.urls] +Homepage = "https://github.com/Agentix-Project/Agentix" + +[project.scripts] +agentix-tito = "agentix.tito.cli:main" + +[tool.uv.sources] +agentixx = { workspace = true } + +# uv_build, like the other plugins: ship under the `agentix.tito` namespace. +[tool.uv.build-backend] +module-name = "agentix.tito" +module-root = "" +namespace = true + +[tool.pytest.ini_options] +testpaths = ["tests"] +pythonpath = ["."] +addopts = "-q" diff --git a/plugins/tito/tests/package/test_cli.py b/plugins/tito/tests/package/test_cli.py new file mode 100644 index 0000000..53d63e6 --- /dev/null +++ b/plugins/tito/tests/package/test_cli.py @@ -0,0 +1,77 @@ +import pytest +from agentix.tito.cli import build_parser, main + + +def test_cli_top_level_help(capsys): + with pytest.raises(SystemExit): + main(["--help"]) + assert "serve" in capsys.readouterr().out + + +def test_cli_serve_help(capsys): + with pytest.raises(SystemExit): + main(["serve", "--help"]) + out = capsys.readouterr().out + assert "--hf-checkpoint" in out + assert "--tito-model" in out + assert "--tito-allowed-append-roles" in out + + +def test_cli_serve_parses_args(): + args = build_parser().parse_args( + [ + "serve", + "--hf-checkpoint", "Qwen/Qwen3-0.6B", + "--backend-url", "http://127.0.0.1:8000", + "--tito-model", "qwen3", + "--tito-allowed-append-roles", "tool", "user", + ] + ) + assert args.command == "serve" + assert args.hf_checkpoint == "Qwen/Qwen3-0.6B" + assert args.tito_model == "qwen3" + assert args.tito_allowed_append_roles == ["tool", "user"] + + +def test_cli_tito_model_choices_are_qwen3_and_default(): + args = build_parser().parse_args(["serve", "--hf-checkpoint", "X"]) + assert args.tito_model == "default" + with pytest.raises(SystemExit): + build_parser().parse_args(["serve", "--hf-checkpoint", "X", "--tito-model", "glm47"]) + + +def test_cli_backend_url_is_repeatable_for_a_pool(): + args = build_parser().parse_args( + ["serve", "--hf-checkpoint", "X", + "--backend-url", "http://r1:8000", "--backend-url", "http://r2:8000", + "--routing-policy", "round_robin"] + ) + assert args.backend_url == ["http://r1:8000", "http://r2:8000"] + assert args.routing_policy == "round_robin" + + +def test_cli_values_plumb_pool_and_trust_flags(): + from agentix.tito.config import TITOGatewayConfig + + cfg = TITOGatewayConfig.from_cli_values( + hf_checkpoint="X", + backend_url=None, + backend_urls=["http://r1:8000", "http://r2:8000"], + routing_policy="round_robin", + trust_remote_code=True, + chat_template_path=None, + tito_model="default", + tito_allowed_append_roles=["tool"], + session_server_ip="127.0.0.1", + session_server_port=30000, + router_timeout=1.0, + ) + assert cfg.backend_urls == ("http://r1:8000", "http://r2:8000") + assert cfg.routing_policy == "round_robin" + assert cfg.trust_remote_code is True + assert cfg.as_session_args().trust_remote_code is True + + +def test_cli_trust_remote_code_defaults_off(): + args = build_parser().parse_args(["serve", "--hf-checkpoint", "X"]) + assert args.trust_remote_code is False diff --git a/plugins/tito/tests/package/test_config_discovery.py b/plugins/tito/tests/package/test_config_discovery.py new file mode 100644 index 0000000..9b6862a --- /dev/null +++ b/plugins/tito/tests/package/test_config_discovery.py @@ -0,0 +1,137 @@ +import pytest +from agentix.tito import discovery +from agentix.tito.config import TITOGatewayConfig + + +def test_explicit_backend_url_wins_over_environment_and_probe(monkeypatch): + env = {"TITO_BACKEND_URL": "http://env.example:3000"} + + def fail_if_probed(*args, **kwargs): + raise AssertionError("explicit backend URL must not probe candidates") + + monkeypatch.setattr(discovery, "probe_backend_url", fail_if_probed) + + assert ( + discovery.discover_backend_url( + "localhost:8000", + env=env, + probe_candidates=("http://probe.example:8000",), + ) + == "http://localhost:8000" + ) + + +def test_environment_precedence_is_deterministic_and_wins_over_probe(monkeypatch): + env = { + "OPENAI_BASE_URL": "http://openai.example:8000", + "SGLANG_BASE_URL": "http://sglang.example:8000", + } + + def fail_if_probed(*args, **kwargs): + raise AssertionError("environment backend URL must not probe candidates") + + monkeypatch.setattr(discovery, "probe_backend_url", fail_if_probed) + + assert ( + discovery.discover_backend_url(env=env, probe_candidates=("http://probe.example:8000",)) + == "http://openai.example:8000" + ) + + +def test_probe_selects_health_success(monkeypatch): + calls = [] + + def endpoint_probe(url, timeout): + calls.append((url, timeout)) + return url == "http://candidate.example:8000/health" + + monkeypatch.setattr(discovery, "_probe_endpoint", endpoint_probe) + + assert ( + discovery.discover_backend_url( + env={}, + probe_candidates=("candidate.example:8000",), + probe_timeout=1.5, + ) + == "http://candidate.example:8000" + ) + assert calls == [("http://candidate.example:8000/health", 1.5)] + + +def test_probe_falls_back_to_models_endpoint(monkeypatch): + calls = [] + + def endpoint_probe(url, timeout): + calls.append(url) + return url == "http://candidate.example:8000/v1/models" + + monkeypatch.setattr(discovery, "_probe_endpoint", endpoint_probe) + + assert ( + discovery.discover_backend_url(env={}, probe_candidates=("http://candidate.example:8000",)) + == "http://candidate.example:8000" + ) + assert calls == [ + "http://candidate.example:8000/health", + "http://candidate.example:8000/v1/models", + ] + + +def test_probe_selection_uses_first_live_candidate(monkeypatch): + calls = [] + + def endpoint_probe(url, timeout): + calls.append(url) + return url == "http://second.example:8000/health" + + monkeypatch.setattr(discovery, "_probe_endpoint", endpoint_probe) + + assert ( + discovery.discover_backend_url( + env={}, + probe_candidates=("http://first.example:8000", "http://second.example:8000"), + ) + == "http://second.example:8000" + ) + assert calls == [ + "http://first.example:8000/health", + "http://first.example:8000/v1/models", + "http://second.example:8000/health", + ] + + +def test_missing_backend_url_fails_clearly(): + with pytest.raises(RuntimeError, match="backend URL not found"): + discovery.discover_backend_url(env={}, probe_candidates=()) + + +def test_no_live_probe_candidate_fails_clearly(monkeypatch): + monkeypatch.setattr(discovery, "_probe_endpoint", lambda url, timeout: False) + + with pytest.raises(RuntimeError, match="start a live backend"): + discovery.discover_backend_url(env={}, probe_candidates=("http://dead.example:8000",)) + + +def test_from_cli_values_maps_fields(): + config = TITOGatewayConfig.from_cli_values( + hf_checkpoint="model", + backend_url="http://backend", + chat_template_path=None, + tito_model="qwen3", + tito_allowed_append_roles=["tool", "user"], + session_server_ip="127.0.0.1", + session_server_port=30000, + router_timeout=30, + backend_probe_candidates=["http://probe-a:8000", "probe-b:8001"], + backend_probe_timeout=2.0, + ) + + assert config.router_timeout == 30 + assert config.tito_allowed_append_roles == ("tool", "user") + assert config.backend_probe_candidates == ("http://probe-a:8000", "probe-b:8001") + assert config.backend_probe_timeout == 2.0 + + +def test_invalid_append_role_fails(): + with pytest.raises(ValueError, match="unsupported tito append roles"): + TITOGatewayConfig(hf_checkpoint="model", tito_allowed_append_roles=("assistant",)) diff --git a/plugins/tito/tests/package/test_engine.py b/plugins/tito/tests/package/test_engine.py new file mode 100644 index 0000000..650bf55 --- /dev/null +++ b/plugins/tito/tests/package/test_engine.py @@ -0,0 +1,294 @@ +"""Self-contained tests for the native TITO engine. + +These build a tiny in-memory tokenizer (no model download) and assert the engine's +invariants directly: the incremental tokenization equals a from-scratch render, the +comparator classifies mismatches correctly, message matching collapses falsy +sentinels, and the session state machine rolls back to the last assistant checkpoint. +""" + +from __future__ import annotations + +import pytest +from agentix.tito.engine.compare import MismatchType, TokenSeqComparator +from agentix.tito.engine.errors import TokenizationError +from agentix.tito.engine.messages import assert_messages_append_only_with_allowed_role, message_matches +from agentix.tito.engine.pretokenize import Qwen3TITOTokenizer, get_tito_tokenizer +from agentix.tito.engine.trajectory import LinearTrajectory, SessionRegistry +from tokenizers import Tokenizer, models, pre_tokenizers +from transformers import PreTrainedTokenizerFast + + +@pytest.fixture(scope="module") +def tok(): + specials = ["", "", "", "<|im_start|>", "<|im_end|>"] + words = ["system", "user", "assistant", "tool", "dummy", "You", "are", "ok", + "done", "compute", "17", "23", "391", "X", "Y", "Hello"] + vocab = {t: i for i, t in enumerate(specials + words)} + tk = Tokenizer(models.WordLevel(vocab=vocab, unk_token="")) + tk.pre_tokenizer = pre_tokenizers.Whitespace() + t = PreTrainedTokenizerFast( + tokenizer_object=tk, unk_token="", bos_token="", eos_token="", + additional_special_tokens=["<|im_start|>", "<|im_end|>"], + ) + t.chat_template = ( + "{%- for m in messages -%}<|im_start|>{{ m['role'] }} {{ m['content'] or '' }}<|im_end|>{%- endfor -%}" + "{%- if add_generation_prompt -%}<|im_start|>assistant {%- endif -%}" + ) + return t + + +def _types(ms): + return [(m.type, m.segment_index) for m in ms] + + +def test_comparator_classifies_mismatches(tok): + cmp = TokenSeqComparator(tok, assistant_start_str="<|im_start|>assistant") + ims, ime = tok.convert_tokens_to_ids("<|im_start|>"), tok.convert_tokens_to_ids("<|im_end|>") + S, U, A = (tok.convert_tokens_to_ids(w) for w in ("system", "user", "assistant")) + Y391, Y23, Yok, YH = (tok.convert_tokens_to_ids(w) for w in ("391", "23", "ok", "Hello")) + + assert cmp.compare_sequences([ims, U, Y391, ime], [ims, U, Y391, ime]) == [] + assert _types(cmp.compare_sequences([ims, U, Y391, ime], [ims, U, Y23, ime])) == [ + (MismatchType.NON_ASSISTANT_TEXT, 1) + ] + assert _types(cmp.compare_sequences([ims, A, Yok, ime], [ims, A, YH, ime])) == [ + (MismatchType.ASSISTANT_TEXT, 1) + ] + assert _types(cmp.compare_sequences([ims, Y391, ime], [ims, Y391, ime, ims])) == [ + (MismatchType.SPECIAL_TOKEN_COUNT, -1) + ] + assert _types(cmp.compare_sequences([ims, Y391, ime], [ime, Y391, ims])) == [ + (MismatchType.SPECIAL_TOKEN_TYPE, 0), + (MismatchType.SPECIAL_TOKEN_TYPE, 2), + ] + # trailing trim removes false structural diffs + assert cmp.compare_sequences([ims, Y391, ime], [ims, Y391, ime, ime], trim_trailing_ids={ime}) == [] + + +def test_message_matches_collapses_falsy_sentinels(): + assert message_matches({"role": "a", "content": ""}, {"role": "a", "content": None}) + assert message_matches({"role": "a", "tool_calls": []}, {"role": "a", "tool_calls": None}) + assert message_matches({"role": "u", "content": "x"}, {"role": "u", "content": "x", "extra": 1}) + # reasoning_content "\n\n" is non-falsy → not collapsed (the bug we hit) + assert not message_matches({"role": "a", "reasoning_content": "\n\n"}, {"role": "a", "reasoning_content": None}) + assert not message_matches({"role": "u", "content": "x"}, {"role": "t", "content": "x"}) + + +def test_append_only_enforced(): + stored = [{"role": "user", "content": "x"}] + assert_messages_append_only_with_allowed_role(stored, stored + [{"role": "tool", "content": "y"}], ["tool"]) + with pytest.raises(ValueError): + assert_messages_append_only_with_allowed_role(stored, stored + [{"role": "user", "content": "z"}], ["tool"]) + with pytest.raises(ValueError): + assert_messages_append_only_with_allowed_role(stored, [{"role": "user", "content": "DIFF"}], ["tool"]) + + +@pytest.mark.parametrize( + "appends", + [ + [{"role": "tool", "content": "391"}], + [{"role": "tool", "content": "391"}, {"role": "tool", "content": "23"}], + [{"role": "user", "content": "Hello"}], + [{"role": "tool", "content": "X"}, {"role": "user", "content": "Y"}], + ], +) +def test_incremental_equals_full_render(tok, appends): + """The core invariant: merge(prefix, incremental) == full from-scratch render.""" + tt = get_tito_tokenizer(tok, "default", allowed_append_roles=("tool", "user")) + old = [{"role": "system", "content": "You are"}, {"role": "user", "content": "compute 17 23"}, + {"role": "assistant", "content": "ok"}] + new = old + appends + prefix = tt.render_messages(old, add_generation_prompt=False, tokenize=True) + merged = tt.merge_tokens(old, new, prefix, None) + full = tt.render_messages(new, add_generation_prompt=True, tokenize=True) + assert merged == full + + +def test_qwen3_newline_fixup(): + class FakeTok: + def encode(self, t, add_special_tokens=False): + return [99] # "\n" -> single id + + def convert_tokens_to_ids(self, t): + return 88 # "<|im_end|>" + + q = Qwen3TITOTokenizer(FakeTok(), chat_template_kwargs={"chat_template": "x"}) + q.tokenize_additional_non_assistant = lambda o, n, t=None: [1, 2, 3] + assert q.merge_tokens([], [], [7, 88], None) == [7, 88, 99, 1, 2, 3] # prefix ends in im_end -> insert \n + assert q.merge_tokens([], [], [7, 5], None) == [7, 5, 1, 2, 3] # otherwise no insert + + +def test_trajectory_rollback_to_assistant_checkpoint(tok): + tt = get_tito_tokenizer(tok, "default", allowed_append_roles=("tool", "user")) + reg = SessionRegistry(None, tok, tito_tokenizer=tt) + tr = LinearTrajectory() + sys = [{"role": "system", "content": "You are"}, {"role": "user", "content": "compute 17 23"}] + a0 = {"role": "assistant", "content": "ok"} + + tr.prepare_pretokenized(sys, None, tito_tokenizer=tt) + tr.update_pretokenized_state( + sys, a0, tt.render_messages(sys + [a0], add_generation_prompt=False, tokenize=True), [], tt.max_trim_tokens + ) + + m1 = sys + [a0, {"role": "tool", "content": "391"}] + a1 = {"role": "assistant", "content": "done"} + tr.prepare_pretokenized(m1, None, tito_tokenizer=tt) + tr.update_pretokenized_state( + m1, a1, tt.render_messages(m1 + [a1], add_generation_prompt=False, tokenize=True), [], tt.max_trim_tokens + ) + assert tr.num_assistant == 2 + assert reg.compute_session_mismatch(tr) == [] # clean chain → no mismatch + + # retry the tool turn with a different result → rollback to a0 checkpoint + tr.prepare_pretokenized(sys + [a0, {"role": "tool", "content": "X"}], None, tito_tokenizer=tt) + assert tr.num_assistant == 1 + assert [m.get("role") for m in tr.messages] == ["system", "user", "assistant"] + + +def test_trajectory_keeps_only_reachable_checkpoints(tok): + """Single-step rollback (MAX_ASSISTANT_ROLLBACK_STEPS=1) can only ever reach + the last two checkpoints; retaining every full prefix+completion token list + is O(turns^2) dead memory per session. Rollback must still work.""" + tt = get_tito_tokenizer(tok, "default", allowed_append_roles=("tool", "user")) + tr = LinearTrajectory() + msgs = [{"role": "system", "content": "You are"}, {"role": "user", "content": "compute 17 23"}] + + for i, reply in enumerate(["ok", "done", "ok", "done"]): + tr.prepare_pretokenized(msgs, None, tito_tokenizer=tt) + asst = {"role": "assistant", "content": reply} + tr.update_pretokenized_state( + msgs, asst, + tt.render_messages(msgs + [asst], add_generation_prompt=False, tokenize=True), + [], tt.max_trim_tokens, + ) + msgs = msgs + [asst, {"role": "tool", "content": "391"}] + + assert tr.num_assistant == 4 + assert len(tr.trajectory_token_ids) <= 2 # dead checkpoints dropped + + # single-step retry (divergent tool result → rollback one assistant) + # still works across the trimmed history + retry = msgs[:-3] + [{"role": "tool", "content": "X"}] + tr.prepare_pretokenized(retry, None, tito_tokenizer=tt) + assert tr.num_assistant == 3 + assert tr.token_ids # the surviving checkpoint is intact + + +def test_prefix_mismatch_diagnostic_handles_shorter_new_sequence(tok): + """A new prompt+completion SHORTER than the stored checkpoint must produce + the TokenizationError diagnostic, not a bare ValueError from zip(strict).""" + tr = LinearTrajectory() + tr.trajectory_token_ids.append(list(range(10))) + tr.num_assistant = 1 + with pytest.raises(TokenizationError): + # prefix-consistent but SHORTER than the stored checkpoint: the + # diagnostic scan finds no differing pair before the short side ends. + tr.update_pretokenized_state( + [{"role": "user", "content": "compute"}], + {"role": "assistant", "content": "ok"}, + prompt_token_ids=[0], + completion_token_ids=[1], + max_trim_tokens=0, + ) + + +def test_rollback_survives_exhausted_checkpoints(tok): + """Retry chains can outrun the trimmed checkpoint window: rollback with no + intervening update (failed proxy), then another legal rollback. The prompt + must then be the full from-scratch render — never a merge onto an empty + prefix that silently drops the whole stored history.""" + tt = get_tito_tokenizer(tok, "default", allowed_append_roles=("tool", "user")) + tr = LinearTrajectory() + msgs = [{"role": "system", "content": "You are"}, {"role": "user", "content": "compute 17 23"}] + for reply in ("ok", "done", "ok"): + tr.prepare_pretokenized(msgs, None, tito_tokenizer=tt) + asst = {"role": "assistant", "content": reply} + tr.update_pretokenized_state( + msgs, asst, + tt.render_messages(msgs + [asst], add_generation_prompt=False, tokenize=True), + [], tt.max_trim_tokens, + ) + msgs = msgs + [asst, {"role": "tool", "content": "391"}] + + assert tr.num_assistant == 3 + + # retry turn 3 (divergent tool after a1) — proxy fails, so no update lands + stored = tr.messages + retry3 = stored[:5] + [{"role": "tool", "content": "X"}] + tr.prepare_pretokenized(retry3, None, tito_tokenizer=tt) + assert tr.num_assistant == 2 + + # retry turn 2 (divergent tool after a0) — checkpoint window exhausted + retry2 = tr.messages[:3] + [{"role": "tool", "content": "Y"}] + prompt = tr.prepare_pretokenized(retry2, None, tito_tokenizer=tt) + assert tr.num_assistant == 1 + assert prompt == tt.render_messages(retry2, add_generation_prompt=True, tokenize=True) + + +def test_rejected_divergent_request_does_not_commit_rollback(tok): + """A divergent request that FAILS validation (disallowed appended role) + must be a pure 400: no rollback side effects may survive, and the original + branch must remain resumable.""" + from agentix.tito.engine.errors import MessageValidationError + + tt = get_tito_tokenizer(tok, "default", allowed_append_roles=("tool",)) + tr = LinearTrajectory() + sys = [{"role": "system", "content": "You are"}, {"role": "user", "content": "compute 17 23"}] + a0 = {"role": "assistant", "content": "ok"} + tr.prepare_pretokenized(sys, None, tito_tokenizer=tt) + tr.update_pretokenized_state( + sys, a0, tt.render_messages(sys + [a0], add_generation_prompt=False, tokenize=True), [], tt.max_trim_tokens + ) + m1 = sys + [a0, {"role": "tool", "content": "391"}] + a1 = {"role": "assistant", "content": "done"} + tr.prepare_pretokenized(m1, None, tito_tokenizer=tt) + tr.update_pretokenized_state( + m1, a1, tt.render_messages(m1 + [a1], add_generation_prompt=False, tokenize=True), [], tt.max_trim_tokens + ) + assert tr.num_assistant == 2 + + # divergent at the tool turn AND appends a disallowed user message + bad = sys + [a0, {"role": "tool", "content": "X"}, {"role": "user", "content": "Hello"}] + with pytest.raises(MessageValidationError): + tr.prepare_pretokenized(bad, None, tito_tokenizer=tt) + assert tr.num_assistant == 2 # rollback was NOT committed + assert len(tr.records) == 0 or tr.num_assistant == 2 # state intact + + # the original branch is still resumable + m2 = m1 + [a1, {"role": "tool", "content": "23"}] + tr.prepare_pretokenized(m2, None, tito_tokenizer=tt) + assert tr.num_assistant == 2 + + +def test_version_advances_on_rollback_and_update(tok): + """`num_assistant` is not monotonic (rollback decrements, update increments) + so it cannot detect concurrent interference; `version` must advance on BOTH.""" + tt = get_tito_tokenizer(tok, "default", allowed_append_roles=("tool", "user")) + tr = LinearTrajectory() + sys = [{"role": "system", "content": "You are"}, {"role": "user", "content": "compute 17 23"}] + a0 = {"role": "assistant", "content": "ok"} + tr.prepare_pretokenized(sys, None, tito_tokenizer=tt) + tr.update_pretokenized_state( + sys, a0, tt.render_messages(sys + [a0], add_generation_prompt=False, tokenize=True), [], tt.max_trim_tokens + ) + m1 = sys + [a0, {"role": "tool", "content": "391"}] + a1 = {"role": "assistant", "content": "done"} + tr.prepare_pretokenized(m1, None, tito_tokenizer=tt) + tr.update_pretokenized_state( + m1, a1, tt.render_messages(m1 + [a1], add_generation_prompt=False, tokenize=True), [], tt.max_trim_tokens + ) + + v0 = tr.version + n0 = tr.num_assistant + # rollback (retry the tool turn) then a fresh update: num_assistant is + # back to n0 but version must have moved. + retry = sys + [a0, {"role": "tool", "content": "X"}] + tr.prepare_pretokenized(retry, None, tito_tokenizer=tt) + a1b = {"role": "assistant", "content": "done"} + tr.update_pretokenized_state( + retry, a1b, + tt.render_messages(retry + [a1b], add_generation_prompt=False, tokenize=True), + [], tt.max_trim_tokens, + ) + assert tr.num_assistant == n0 + assert tr.version != v0 diff --git a/plugins/tito/tests/package/test_import_surface.py b/plugins/tito/tests/package/test_import_surface.py new file mode 100644 index 0000000..820f0dc --- /dev/null +++ b/plugins/tito/tests/package/test_import_surface.py @@ -0,0 +1,40 @@ +import pytest + + +def test_public_import_surface(): + import agentix.tito + from agentix.tito import SessionServer, TITOGateway, TITOGatewayConfig, get_tito_tokenizer + + assert agentix.tito.TITOGateway is TITOGateway + assert agentix.tito.TITOGatewayConfig is TITOGatewayConfig + assert agentix.tito.SessionServer is SessionServer + assert callable(get_tito_tokenizer) + + +def test_config_requires_hf_checkpoint(): + from agentix.tito import TITOGatewayConfig + + with pytest.raises(ValueError, match="hf_checkpoint is required"): + TITOGatewayConfig(hf_checkpoint="") + + +def test_gateway_constructs_with_explicit_backend(monkeypatch): + import agentix.tito.gateway as gateway_module + from agentix.tito import TITOGateway + + class FakeSessionServer: + def __init__(self, args, backend_url): + from fastapi import FastAPI + + self.args = args + self.backend_url = backend_url + # A real app: the gateway registers a `/healthz` alias on it at construct. + self.app = FastAPI() + + monkeypatch.setattr(gateway_module, "SessionServer", FakeSessionServer) + + gateway = TITOGateway.from_server(hf_checkpoint="Qwen/Qwen3-0.6B", backend_url="127.0.0.1:8000") + + assert gateway.config.backend_url == "http://127.0.0.1:8000" + assert gateway.app is gateway.server.app + assert gateway.server.args.hf_checkpoint == "Qwen/Qwen3-0.6B" diff --git a/plugins/tito/tests/test_gateway_http.py b/plugins/tito/tests/test_gateway_http.py new file mode 100644 index 0000000..e4d8948 --- /dev/null +++ b/plugins/tito/tests/test_gateway_http.py @@ -0,0 +1,393 @@ +"""HTTP-surface tests: drive the REAL gateway FastAPI app over ASGI. + +Unlike the unit tests (which call `do_proxy` / pool methods directly), these +exercise the full route stack — request parsing, the session lock phases, the +forget-on-delete middleware, error handlers — via real HTTP requests through +`httpx.ASGITransport`. The upstream replica is an `httpx.MockTransport` +swapped into the pooled backend: an OpenAI-shaped chat-completions endpoint +with the `meta_info.output_token_logprobs` the TITO flow requires, keyed by +host so multi-backend routing is observable. The tokenizer is the same tiny +in-memory WordLevel one the engine tests use (no model download). +""" + +from __future__ import annotations + +import json +import types + +import httpx +import pytest +from agentix.tito.pool import BackendPool +from agentix.tito.server import SessionServer +from tokenizers import Tokenizer, models, pre_tokenizers +from transformers import PreTrainedTokenizerFast + +A, B = "http://replica-a:8000", "http://replica-b:8000" + + +@pytest.fixture(scope="module") +def tok(): + specials = ["", "", "", "<|im_start|>", "<|im_end|>"] + words = ["system", "user", "assistant", "tool", "You", "are", "ok", "done", "Hello"] + vocab = {t: i for i, t in enumerate(specials + words)} + tk = Tokenizer(models.WordLevel(vocab=vocab, unk_token="")) + tk.pre_tokenizer = pre_tokenizers.Whitespace() + t = PreTrainedTokenizerFast( + tokenizer_object=tk, unk_token="", bos_token="", eos_token="", + additional_special_tokens=["<|im_start|>", "<|im_end|>"], + ) + t.chat_template = ( + "{%- for m in messages -%}<|im_start|>{{ m['role'] }} {{ m['content'] or '' }}<|im_end|>{%- endfor -%}" + "{%- if add_generation_prompt -%}<|im_start|>assistant {%- endif -%}" + ) + return t + + +def _args(): + return types.SimpleNamespace( + hf_checkpoint="tiny-in-memory", + chat_template_path=None, + tito_allowed_append_roles=None, + tito_model="default", + session_server_instance_id=None, + router_timeout=5.0, + ) + + +class _Replica: + """OpenAI-shaped fake replica behind an httpx.MockTransport. Records every + body it sees; hosts in `down` raise a transport error instead of answering. + `message` / `usage` / `raw_json` allow per-test response shaping.""" + + def __init__(self) -> None: + self.calls: list[tuple[str, dict]] = [] + self.down: set[str] = set() + self.message: dict = {"role": "assistant", "content": "ok done"} + self.usage: dict = {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5} + self.raw_json: dict | None = None # full-body override (still 200) + self.raw_content: bytes | None = None # verbatim body override (still 200) + + def handler(self, request: httpx.Request) -> httpx.Response: + host = request.url.host + if host in self.down: + raise httpx.ConnectError(f"{host} refused") + if request.method == "DELETE": + return httpx.Response(204) + body = json.loads(request.content) + self.calls.append((host, body)) + if self.raw_content is not None: + return httpx.Response( + 200, content=self.raw_content, headers={"content-type": "application/json"} + ) + if self.raw_json is not None: + return httpx.Response(200, json=self.raw_json) + completion_ids = [7, 8] # "ok done" in the tiny vocab + return httpx.Response( + 200, + json={ + "id": "c1", "object": "chat.completion", "model": "m", + "choices": [{ + "index": 0, + "finish_reason": "stop", + "message": self.message, + "meta_info": { + "output_token_logprobs": [[-0.1, t, ""] for t in completion_ids], + "completion_tokens": len(completion_ids), + }, + }], + "usage": self.usage, + }, + ) + + +@pytest.fixture() +def server(tok, monkeypatch): + """(SessionServer over the tiny tokenizer, replica, pool) with a 2-replica pool.""" + monkeypatch.setattr( + "agentix.tito.engine.session_app.load_tokenizer", lambda *a, **k: tok + ) + pool = BackendPool([A, B], policy="sticky", down_cooldown=0.05) + srv = SessionServer(_args(), pool) + replica = _Replica() + srv._backend.client = httpx.AsyncClient( + transport=httpx.MockTransport(replica.handler), timeout=5.0 + ) + return srv, replica, pool + + +@pytest.fixture() +def gateway(server): + """(http client for the real app, replica, pool).""" + srv, replica, pool = server + client = httpx.AsyncClient( + transport=httpx.ASGITransport(app=srv.app), base_url="http://gw", timeout=5.0 + ) + return client, replica, pool + + +_CHAT = {"model": "m", "messages": [{"role": "user", "content": "Hello"}]} + + +@pytest.mark.asyncio +async def test_full_session_flow_over_http(gateway): + client, replica, _ = gateway + r = await client.post("/sessions") + assert r.status_code == 200 + sid = r.json()["session_id"] + + r = await client.post(f"/sessions/{sid}/v1/chat/completions", json=_CHAT) + assert r.status_code == 200 + assert r.json()["choices"][0]["message"]["content"] == "ok done" + # the gateway injected pretokenized input_ids + forced logprobs upstream + _, seen = replica.calls[0] + assert seen["logprobs"] is True + assert isinstance(seen["input_ids"], list) and seen["input_ids"] + + r = await client.get(f"/sessions/{sid}") + assert r.status_code == 200 + got = r.json() + assert len(got["records"]) == 1 + assert got["metadata"]["accumulated_token_ids"][-2:] == [7, 8] + + r = await client.delete(f"/sessions/{sid}") + assert r.status_code == 204 + r = await client.get(f"/sessions/{sid}") + assert r.status_code == 404 + + +@pytest.mark.asyncio +async def test_stream_request_is_forced_non_streaming(gateway): + """The TITO flow needs the full JSON completion (logprobs + meta_info), so + the gateway must force stream=false upstream and answer 200 with the JSON + body — not 500 on an unparseable SSE stream.""" + client, replica, _ = gateway + sid = (await client.post("/sessions")).json()["session_id"] + + r = await client.post( + f"/sessions/{sid}/v1/chat/completions", json={**_CHAT, "stream": True} + ) + assert r.status_code == 200 + assert r.json()["choices"][0]["message"]["content"] == "ok done" + _, seen = replica.calls[0] + assert seen["stream"] is False + + +@pytest.mark.asyncio +async def test_tool_call_completion_with_null_content_is_accepted(gateway): + """Tool-call-only assistant turns routinely carry content:null (sglang's + parser consumes all text; vLLM emits None) — the gateway must record the + turn, not 502 an agentic rollout at its first tool call.""" + client, replica, _ = gateway + replica.message = { + "role": "assistant", + "content": None, + "tool_calls": [{ + "id": "call_1", "type": "function", + "function": {"name": "compute", "arguments": "{}"}, + }], + } + sid = (await client.post("/sessions")).json()["session_id"] + r = await client.post(f"/sessions/{sid}/v1/chat/completions", json=_CHAT) + assert r.status_code == 200 + assert r.json()["choices"][0]["message"]["tool_calls"][0]["id"] == "call_1" + got = (await client.get(f"/sessions/{sid}")).json() + assert len(got["records"]) == 1 # the turn was recorded + + +@pytest.mark.asyncio +async def test_content_none_without_tool_calls_is_still_502(gateway): + client, replica, _ = gateway + replica.message = {"role": "assistant", "content": None} + sid = (await client.post("/sessions")).json()["session_id"] + r = await client.post(f"/sessions/{sid}/v1/chat/completions", json=_CHAT) + assert r.status_code == 502 + + +@pytest.mark.asyncio +async def test_malformed_request_body_is_400(gateway): + client, _, _ = gateway + sid = (await client.post("/sessions")).json()["session_id"] + r = await client.post( + f"/sessions/{sid}/v1/chat/completions", + content=b"{not json", + headers={"content-type": "application/json"}, + ) + assert r.status_code == 400 + + +@pytest.mark.asyncio +async def test_malformed_200_upstream_is_clean_502(gateway): + """A 200 upstream body without the expected structure (choices/meta_info + shapes) must surface as a clean 502, not an unhandled 500.""" + client, replica, _ = gateway + sid = (await client.post("/sessions")).json()["session_id"] + for weird in ({"weird": True}, {"choices": []}, {"choices": [{"meta_info": None}]}): + replica.raw_json = weird + r = await client.post(f"/sessions/{sid}/v1/chat/completions", json=_CHAT) + assert r.status_code == 502, weird + + +@pytest.mark.asyncio +async def test_upstream_nan_json_passes_through(gateway): + """The proxy must not re-encode the upstream JSON body — NaN/Infinity are + valid for Python's json but not for a strict re-serializer, and re-encoding + also perturbs key order/whitespace for byte-sensitive consumers.""" + client, replica, _ = gateway + # stdlib json accepts NaN both ways; httpx's own json= encoder does not, + # so the replica ships the body verbatim. + replica.raw_content = json.dumps({ + "id": "c1", "object": "chat.completion", "model": "m", + "choices": [{ + "index": 0, "finish_reason": "stop", + "message": {"role": "assistant", "content": "ok done"}, + "meta_info": { + "output_token_logprobs": [[float("nan"), 7, ""], [-0.1, 8, ""]], + "completion_tokens": 2, + }, + }], + "usage": {}, + }).encode() + sid = (await client.post("/sessions")).json()["session_id"] + r = await client.post(f"/sessions/{sid}/v1/chat/completions", json=_CHAT) + assert r.status_code == 200 + assert b"NaN" in r.content + + +@pytest.mark.asyncio +async def test_registry_tokenizer_load_is_not_trust_remote_code(tok, monkeypatch): + """trust_remote_code executes checkpoint-shipped Python at startup — it + must be an explicit opt-in, never the hardcoded default.""" + seen: dict = {} + + def fake_load(name, chat_template_path=None, trust_remote_code=False): + seen["trust_remote_code"] = trust_remote_code + return tok + + monkeypatch.setattr("agentix.tito.engine.session_app.load_tokenizer", fake_load) + pool = BackendPool([A]) + SessionServer(_args(), pool) + assert seen["trust_remote_code"] is False + + +@pytest.mark.asyncio +async def test_unknown_session_is_404_over_http(gateway): + client, _, _ = gateway + r = await client.post("/sessions/nope/v1/chat/completions", json=_CHAT) + assert r.status_code == 404 + + +@pytest.mark.asyncio +async def test_backend_failover_and_recovery_over_http(gateway): + """A replica that refuses connections is marked down and the session is + repinned to the survivor; once the cooldown passes the failed replica is + eligible again for new sessions.""" + import asyncio + + client, replica, pool = gateway + sid = (await client.post("/sessions")).json()["session_id"] + pinned = pool.pick(sid) + + replica.down.add(httpx.URL(pinned).host) + r = await client.post(f"/sessions/{sid}/v1/chat/completions", json=_CHAT) + assert r.status_code == 502 # transport error surfaces as 502, marks down + + r = await client.post(f"/sessions/{sid}/v1/chat/completions", json=_CHAT) + assert r.status_code == 200 # repinned to the healthy replica + survivor = replica.calls[-1][0] + assert survivor != httpx.URL(pinned).host + + replica.down.clear() + await asyncio.sleep(0.06) # let the down cooldown expire + hosts = {pool.pick(f"fresh-{i}") for i in range(4)} + assert pinned in hosts # recovered replica takes new sessions again + + +@pytest.mark.asyncio +async def test_delete_forgets_pin_over_http(gateway): + client, _, pool = gateway + sid = (await client.post("/sessions")).json()["session_id"] + await client.post(f"/sessions/{sid}/v1/chat/completions", json=_CHAT) + assert sid in pool._assigned # noqa: SLF001 + r = await client.delete(f"/sessions/{sid}") + assert r.status_code == 204 + assert sid not in pool._assigned # noqa: SLF001 + + +@pytest.mark.asyncio +async def test_proxied_subresource_delete_keeps_pin(gateway): + """Only `DELETE /sessions/{sid}` ends the session; a DELETE proxied to a + backend sub-resource under the session prefix must not drop the sticky pin.""" + client, _, pool = gateway + sid = (await client.post("/sessions")).json()["session_id"] + await client.post(f"/sessions/{sid}/v1/chat/completions", json=_CHAT) + assert sid in pool._assigned # noqa: SLF001 + r = await client.delete(f"/sessions/{sid}/v1/some/backend/resource") + assert r.status_code == 204 + assert sid in pool._assigned # noqa: SLF001 + + +@pytest.mark.asyncio +async def test_slow_request_body_does_not_hold_session_lock(server): + """The session lock must not be held while reading the request body — a + dribbling client upload would otherwise pin the lock indefinitely, wedging + DELETE and every other operation on the session.""" + import asyncio + + srv, _, _ = server + client = httpx.AsyncClient( + transport=httpx.ASGITransport(app=srv.app), base_url="http://gw", timeout=5.0 + ) + sid = (await client.post("/sessions")).json()["session_id"] + session = srv.app.state.tito_registry.get_session(sid) + + payload = json.dumps(_CHAT).encode() + started = asyncio.Event() + release = asyncio.Event() + + async def dribble(): + yield payload[:5] + started.set() + await release.wait() + yield payload[5:] + + task = asyncio.create_task(client.post( + f"/sessions/{sid}/v1/chat/completions", + content=dribble(), + headers={"content-type": "application/json"}, + )) + await started.wait() + await asyncio.sleep(0.01) # let the handler park inside the body read + locked_during_upload = session.lock.locked() + release.set() + r = await task + assert r.status_code == 200 + assert locked_during_upload is False + + +@pytest.mark.asyncio +async def test_cancelled_delete_does_not_brick_session(server): + """A DELETE cancelled while waiting for the session lock (client timeout / + disconnect) must leave the session deletable — not wedged behind a stuck + closing flag with the trajectory leaked in the registry.""" + import asyncio + import contextlib + + srv, _, _ = server + client = httpx.AsyncClient( + transport=httpx.ASGITransport(app=srv.app), base_url="http://gw", timeout=5.0 + ) + sid = (await client.post("/sessions")).json()["session_id"] + session = srv.app.state.tito_registry.get_session(sid) + + await session.lock.acquire() # simulate another request holding the lock + try: + del_task = asyncio.create_task(client.delete(f"/sessions/{sid}")) + await asyncio.sleep(0.05) # let the DELETE park at lock acquisition + del_task.cancel() + with contextlib.suppress(asyncio.CancelledError, httpx.HTTPError): + await del_task + finally: + session.lock.release() + + r = await client.delete(f"/sessions/{sid}") + assert r.status_code == 204 diff --git a/plugins/tito/tests/test_pool.py b/plugins/tito/tests/test_pool.py new file mode 100644 index 0000000..5a36e07 --- /dev/null +++ b/plugins/tito/tests/test_pool.py @@ -0,0 +1,105 @@ +"""Unit tests for the Gateway backend pool routing (no model in the loop).""" + +from __future__ import annotations + +import time + +import pytest +from agentix.tito.pool import BackendPool + +A, B, C = "http://h1:8000", "http://h2:8000", "http://h3:8000" + + +def test_requires_backends() -> None: + with pytest.raises(ValueError): + BackendPool([]) + + +def test_bad_policy() -> None: + with pytest.raises(ValueError): + BackendPool([A], policy="nope") + + +def test_single_backend_always() -> None: + pool = BackendPool([A]) + assert pool.pick("s1") == A + assert pool.pick() == A + + +def test_sticky_pins_session_to_one_backend() -> None: + pool = BackendPool([A, B, C], policy="sticky") + first = pool.pick("rollout-1") + # Same session keeps hitting the same backend across many turns. + assert all(pool.pick("rollout-1") == first for _ in range(10)) + + +def test_sticky_spreads_distinct_sessions_round_robin() -> None: + pool = BackendPool([A, B, C], policy="sticky") + assigned = [pool.pick(f"s{i}") for i in range(3)] + assert sorted(assigned) == sorted([A, B, C]) # 3 sessions → 3 distinct backends + + +def test_round_robin_cycles_every_request() -> None: + pool = BackendPool([A, B], policy="round_robin") + assert [pool.pick("ignored") for _ in range(4)] == [A, B, A, B] + + +def test_down_backend_is_skipped() -> None: + pool = BackendPool([A, B], policy="round_robin") + pool.report_down(A) + assert {pool.pick() for _ in range(6)} == {B} + pool.report_up(A) + assert A in {pool.pick() for _ in range(6)} + + +def test_sticky_session_reassigned_when_backend_down() -> None: + pool = BackendPool([A, B], policy="sticky") + pinned = pool.pick("r1") + pool.report_down(pinned) + reassigned = pool.pick("r1") + assert reassigned != pinned + assert reassigned not in pool._down + + +def test_all_down_falls_back_not_fails() -> None: + pool = BackendPool([A, B], policy="round_robin") + pool.report_down(A) + pool.report_down(B) + # Better to attempt a (maybe-recovered) backend than fail routing outright. + assert pool.pick() in (A, B) + + +def test_down_backend_retried_after_cooldown() -> None: + """Nothing in the serving path calls `report_up` for a backend that is + never picked — without a cooldown a downed replica is blacklisted forever. + After `down_cooldown` the backend must become eligible again (half-open).""" + pool = BackendPool([A, B], policy="round_robin", down_cooldown=0.05) + pool.report_down(A) + assert {pool.pick() for _ in range(4)} == {B} + time.sleep(0.06) + assert A in {pool.pick() for _ in range(4)} + + +def test_sticky_pin_kept_when_cooldown_expires() -> None: + """Once the pinned backend's cooldown expires it is retryable — the session + should return to its original pin (prefix cache) rather than stay migrated.""" + pool = BackendPool([A, B], policy="sticky", down_cooldown=0.05) + pinned = pool.pick("r1") + pool.report_down(pinned) + time.sleep(0.06) + assert pool.pick("r1") == pinned + + +def test_sticky_pins_are_bounded_lru() -> None: + """Session pins are only dropped by an explicit DELETE, but the recommended + rollout flow never deletes (the trajectory must survive for harvest) — the + pin map must be bounded, evicting the least-recently-used session.""" + pool = BackendPool([A, B], policy="sticky", max_pins=4) + pins = {f"s{i}": pool.pick(f"s{i}") for i in range(10)} + assert len(pool._assigned) == 4 # noqa: SLF001 + # the most recently seen sessions keep their pins ... + for sid in ("s6", "s7", "s8", "s9"): + assert pool.pick(sid) == pins[sid] + # ... and the oldest were evicted (repinning is allowed — stickiness is a + # cache-locality optimization, not a correctness requirement) + assert "s0" not in pool._assigned # noqa: SLF001 diff --git a/plugins/tito/tests/test_pool_routing.py b/plugins/tito/tests/test_pool_routing.py new file mode 100644 index 0000000..e6f9af6 --- /dev/null +++ b/plugins/tito/tests/test_pool_routing.py @@ -0,0 +1,117 @@ +"""Wiring tests for BackendPool routing in the SessionServer (no model/GPU). + +Uses ``hf_checkpoint=None`` so the session server skips tokenizer/route setup — +we drive the pool-aware ``do_proxy`` / forget hook directly. +""" + +from __future__ import annotations + +import types + +import pytest +from agentix.tito.pool import BackendPool +from agentix.tito.server import SessionServer, _session_id_from_path + +A = "http://a:8000" +B = "http://b:8000" + + +def _args(): + return types.SimpleNamespace(hf_checkpoint=None, router_timeout=600.0) + + +class _URL: + def __init__(self, path: str, query: str = "") -> None: + self.path = path + self.query = query + + +class _Request: + def __init__(self, path: str, method: str = "POST", body: bytes = b"{}") -> None: + self.url = _URL(path) + self.method = method + self.headers = {} + self._body = body + + async def body(self) -> bytes: + return self._body + + +class _Resp: + def __init__(self, status: int = 200) -> None: + self.status_code = status + self.headers = {} + + async def aread(self) -> bytes: + return b"{}" + + +def test_session_id_from_path(): + assert _session_id_from_path("/sessions/abc/v1/chat/completions") == "abc" + assert _session_id_from_path("/sessions/xyz") == "xyz" + assert _session_id_from_path("/health") is None + assert _session_id_from_path("/") is None + + +@pytest.mark.asyncio +async def test_sticky_routing_pins_session(monkeypatch): + pool = BackendPool([A, B], policy="sticky") + srv = SessionServer(_args(), pool) + seen: list[str] = [] + + async def fake_request(method, url, content=None, headers=None): + seen.append(url) + return _Resp() + + monkeypatch.setattr(srv._backend.client, "request", fake_request) + for _ in range(3): + await srv._backend.do_proxy(_Request("/sessions/s1/v1/chat/completions"), "v1/chat/completions") + # all three turns of one session hit the same backend (prefix-cache locality) + assert len({u.split("/v1/")[0] for u in seen}) == 1 + + +@pytest.mark.asyncio +async def test_transport_error_reports_backend_down(monkeypatch): + import httpx + + pool = BackendPool([A, B], policy="sticky") + srv = SessionServer(_args(), pool) + + async def boom(method, url, content=None, headers=None): + raise httpx.ConnectError("refused") + + monkeypatch.setattr(srv._backend.client, "request", boom) + result = await srv._backend.do_proxy(_Request("/sessions/s9/v1/chat/completions"), "v1/chat/completions") + assert result["status_code"] == 502 + # the picked backend was marked down + assert pool._down # noqa: SLF001 - asserting routing side effect + + +@pytest.mark.asyncio +async def test_successful_proxy_reports_backend_up(monkeypatch): + """A backend that answers a proxied request is healthy — the success path + must clear a stale down mark (the only other recovery is the cooldown).""" + pool = BackendPool([A], policy="sticky") + pool.report_down(A) # sole backend: the all-down fallback still picks it + srv = SessionServer(_args(), pool) + + async def ok(method, url, content=None, headers=None): + return _Resp() + + monkeypatch.setattr(srv._backend.client, "request", ok) + await srv._backend.do_proxy(_Request("/sessions/s1/v1/chat/completions"), "v1/chat/completions") + assert not pool._down # noqa: SLF001 + + +@pytest.mark.asyncio +async def test_forget_on_delete_drops_pin(): + pool = BackendPool([A, B], policy="sticky") + pool.pick("s2") + assert "s2" in pool._assigned # noqa: SLF001 + srv = SessionServer(_args(), pool) + + async def call_next(_req): + return _Resp(status=204) + + await srv._forget_on_delete(_Request("/sessions/s2", method="DELETE"), call_next) + assert "s2" not in pool._assigned # noqa: SLF001 diff --git a/pyproject.toml b/pyproject.toml index fec08c4..b9bd318 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -84,6 +84,7 @@ members = [ "plugins/abridge", "plugins/runner", "plugins/runtime-basic", + "plugins/tito", "plugins/trace-otel", "plugins/agents/*", "plugins/datasets/*", @@ -122,6 +123,7 @@ include = [ "plugins/providers/daytona/agentix", "plugins/providers/e2b/agentix", "plugins/runtime-basic/agentix", + "plugins/tito/agentix", "plugins/trace-otel/agentix", ] exclude = ["**/__pycache__", "**/.venv", ".venv", "**/build"] @@ -143,6 +145,7 @@ extraPaths = [ "plugins/providers/daytona", "plugins/providers/e2b", "plugins/runtime-basic", + "plugins/tito", "plugins/trace-otel", ] typeCheckingMode = "basic" diff --git a/uv.lock b/uv.lock index fa4ff2c..86d4313 100644 --- a/uv.lock +++ b/uv.lock @@ -29,6 +29,7 @@ members = [ "agentix-provider-e2b", "agentix-runner", "agentix-runtime-basic", + "agentix-tito", "agentix-trace-otel", "agentixx", ] @@ -193,6 +194,45 @@ dependencies = [ [package.metadata] requires-dist = [{ name = "agentixx", editable = "." }] +[[package]] +name = "agentix-tito" +version = "0.1.0" +source = { editable = "plugins/tito" } +dependencies = [ + { name = "agentixx" }, + { name = "fastapi" }, + { name = "httpx" }, + { name = "huggingface-hub" }, + { name = "jinja2" }, + { name = "pydantic", version = "2.11.10", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' and python_full_version < '3.14'" }, + { name = "pydantic", version = "2.13.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12' or python_full_version >= '3.14'" }, + { name = "tokenizers" }, + { name = "transformers" }, + { name = "uvicorn" }, +] + +[package.optional-dependencies] +test = [ + { name = "pytest" }, + { name = "pytest-asyncio" }, +] + +[package.metadata] +requires-dist = [ + { name = "agentixx", editable = "." }, + { name = "fastapi", specifier = ">=0.110" }, + { name = "httpx", specifier = ">=0.27" }, + { name = "huggingface-hub", specifier = ">=0.23" }, + { name = "jinja2", specifier = ">=3.1" }, + { name = "pydantic", specifier = ">=2" }, + { name = "pytest", marker = "extra == 'test'", specifier = ">=8" }, + { name = "pytest-asyncio", marker = "extra == 'test'", specifier = ">=0.23" }, + { name = "tokenizers", specifier = ">=0.19" }, + { name = "transformers", specifier = ">=4.44" }, + { name = "uvicorn", specifier = ">=0.29" }, +] +provides-extras = ["test"] + [[package]] name = "agentix-trace-otel" version = "0.1.0" @@ -3291,6 +3331,30 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/9b/36/9c015cd052fca743dae8cb2aeb16b551444787467db42ceab0fc968865af/ruff-0.15.13-py3-none-win_arm64.whl", hash = "sha256:2471da9bd1068c8c064b5fd9c0c4b6dddffd6369cb1cd68b29993b1709ff1b21", size = 11179336, upload-time = "2026-05-14T13:44:33.026Z" }, ] +[[package]] +name = "safetensors" +version = "0.8.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/45/06/f955dbbb1859e3bd23c8ac6141af5106e7ad5fedec4a3a6e3d60f94b7001/safetensors-0.8.0.tar.gz", hash = "sha256:fabaf3e0f18a6618d9b36560682562157f77c2b71fcffc7b432be2baed9d753d", size = 325846, upload-time = "2026-06-09T07:52:25.563Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/39/a0/f718cda65b05407d228f97602cf60dca269c979867aa5beb25410de26cd3/safetensors-0.8.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:c554f85858e05226d3c2828e32395e677434685d6d94594a41643361c5e837f0", size = 473568, upload-time = "2026-06-09T07:52:18.829Z" }, + { url = "https://files.pythonhosted.org/packages/f5/b1/fa7c600e7dceae12e9606c7578cbc9ff1e1ed55844883ee5c92205e86226/safetensors-0.8.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:c80201d22cbf405b80647a60ada77bba06c8fba2da2743ba1e89cdcc39a81f25", size = 484562, upload-time = "2026-06-09T07:52:17.518Z" }, + { url = "https://files.pythonhosted.org/packages/09/7d/65a7de0af421317bb36a067241e4235fff194eed60b961ed6d3f59a3fc60/safetensors-0.8.0-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7a46e5ff292c356d6991e60942ba7f79817682d3a2cef0702136448cb9c4d235", size = 502844, upload-time = "2026-06-09T07:52:07.624Z" }, + { url = "https://files.pythonhosted.org/packages/91/4f/3175c9d75634e0e0dda0082794193521035edd7c70a6f212bf33ca06ddf4/safetensors-0.8.0-cp310-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:4124502b78f03534117c848f87a39b8f31e577b15eff423bf8bfb95f2a8c30d0", size = 511823, upload-time = "2026-06-09T07:52:09.565Z" }, + { url = "https://files.pythonhosted.org/packages/20/87/846c289e7aa2299eff406335717cf43ce8777194ece8aad75772e0411615/safetensors-0.8.0-cp310-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:7bc0a787ba8a35be368ee3574edfa2b1ad389eebd0a72e482ae275490e3f6c98", size = 633461, upload-time = "2026-06-09T07:52:11.128Z" }, + { url = "https://files.pythonhosted.org/packages/76/22/8d64d9df2c45d5ded401df889d0ad90882804ca172d79ec4f0df8f727fe0/safetensors-0.8.0-cp310-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:040070828e36dc8e122178bbbd5830ff9e97920affb84cbe0f46442497bed358", size = 545148, upload-time = "2026-06-09T07:52:13.603Z" }, + { url = "https://files.pythonhosted.org/packages/28/50/f203ff3a3ddfe19308efc83c5a3a29ed02bf786732ec35e68bf9162f3365/safetensors-0.8.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fd6f3f93c9a0a7cc2788ee63fb763353d4bd2e89b0751bc78fcf7dda00bea774", size = 516040, upload-time = "2026-06-09T07:52:16.29Z" }, + { url = "https://files.pythonhosted.org/packages/46/fb/cdaed17ceb2948784fd9c36b6fd3e951b608547cea81a48e8ee6f8cfdfcb/safetensors-0.8.0-cp310-abi3-manylinux_2_31_riscv64.whl", hash = "sha256:fcdd41ec4628fee5799f807c73c353629130fbd942aa23d83c623dd6c9d52d78", size = 513832, upload-time = "2026-06-09T07:52:12.37Z" }, + { url = "https://files.pythonhosted.org/packages/0d/49/1e15de264dcc3b77943d2d0c56a95809956883b1c2d6d585c792523f180b/safetensors-0.8.0-cp310-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:8e9f537aa183a38ace122d27303dcd986b26bd2a7591f9181d7f0c396f4677ca", size = 559930, upload-time = "2026-06-09T07:52:14.743Z" }, + { url = "https://files.pythonhosted.org/packages/2a/43/bf38443278eab4b1be1fce2931e2b012ad9cb7df52ada751d0aab8f7659a/safetensors-0.8.0-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:87eec7ffed2b809f05a398a8becb7d013f19f7837cd15d9748580d6cf30dbaf4", size = 678670, upload-time = "2026-06-09T07:52:20.032Z" }, + { url = "https://files.pythonhosted.org/packages/72/e3/68cd3fa5b48488e84add63e04cb12f3bc28ae4638c06d4508c6e88823d0e/safetensors-0.8.0-cp310-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:4a95ae2b05d7726d751da4ebf626a2ca782b706e101bd894c95bc2450b1cffcc", size = 786679, upload-time = "2026-06-09T07:52:21.322Z" }, + { url = "https://files.pythonhosted.org/packages/29/4b/1c19c509d56e01f4fbb3d0a2e597450f6cc04d1d56cf52defb0a62dfd715/safetensors-0.8.0-cp310-abi3-musllinux_1_2_i686.whl", hash = "sha256:3ae091f16662658bdc019a4ff6cb4c085bb7d725eb5978b183ffd265863b6d2d", size = 765683, upload-time = "2026-06-09T07:52:22.594Z" }, + { url = "https://files.pythonhosted.org/packages/27/43/41c1621732edd934d868a00d1b891584c892a7b62a9aab82ea5a0a5623ee/safetensors-0.8.0-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:8e080062fcde23be189565e1c3305d16751a218ecf9412c8601e64204eb6f846", size = 722361, upload-time = "2026-06-09T07:52:23.924Z" }, + { url = "https://files.pythonhosted.org/packages/8e/3f/73ccf82579412b4a71c4ca673f10b5f1f888d7cf5af7fe24f27d30307be4/safetensors-0.8.0-cp310-abi3-win32.whl", hash = "sha256:2ddf52eac562eda224f99acfa7889d02968c1fd59a5b011ae7d8137c37e9c02d", size = 342401, upload-time = "2026-06-09T07:52:28.895Z" }, + { url = "https://files.pythonhosted.org/packages/1b/6d/3fba214c1e5e0f69991677ec3bc17023f0421776975e1de0c682dca475e2/safetensors-0.8.0-cp310-abi3-win_amd64.whl", hash = "sha256:096ec1a98435df7beb08853bb5aa9081a84f23d0adc67ed1a0a10550f608373f", size = 355540, upload-time = "2026-06-09T07:52:27.832Z" }, + { url = "https://files.pythonhosted.org/packages/8d/fc/7eedc3510d97878876e32774eebbeb61c43f148a96e915c84229a3e967aa/safetensors-0.8.0-cp310-abi3-win_arm64.whl", hash = "sha256:f7838e5135a406ad3e02efdcb8cf2e5397d368b0154537c4fec682dbc544d452", size = 340500, upload-time = "2026-06-09T07:52:26.745Z" }, +] + [[package]] name = "shellingham" version = "1.5.4" @@ -3484,29 +3548,28 @@ wheels = [ [[package]] name = "tokenizers" -version = "0.23.1" +version = "0.22.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "huggingface-hub" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/c1/60/21f715d9faba5f5407ff759472ade058ec4a507ad62bcea47cb847239a73/tokenizers-0.23.1.tar.gz", hash = "sha256:1feeeadf865a7915adc25445dea30e9933e593c31bb96c277cee36de227c8bfa", size = 365748, upload-time = "2026-04-27T14:43:25.606Z" } +sdist = { url = "https://files.pythonhosted.org/packages/73/6f/f80cfef4a312e1fb34baf7d85c72d4411afde10978d4657f8cdd811d3ccc/tokenizers-0.22.2.tar.gz", hash = "sha256:473b83b915e547aa366d1eee11806deaf419e17be16310ac0a14077f1e28f917", size = 372115, upload-time = "2026-01-05T10:45:15.988Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/87/39/b87a87d5bb9470610b80a2d31df42fcffeaf35118b8b97952b2aff598cc7/tokenizers-0.23.1-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:e03d6ffcbe0d56ee9c1ccd070e70a13fa750727c0277e138152acbc0252c2224", size = 3146732, upload-time = "2026-04-27T14:43:15.427Z" }, - { url = "https://files.pythonhosted.org/packages/e2/6a/068ed9f6e444c9d7e9d55ce134181325700f3d7f30410721bdc8f848d727/tokenizers-0.23.1-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:e0948bbb1ac1d7cdfc9fb6d62c596e3b7550036ad60ecd654a66ad273326324e", size = 3054954, upload-time = "2026-04-27T14:43:13.745Z" }, - { url = "https://files.pythonhosted.org/packages/6c/36/e006edf031154cba92b8416057d92c3abe3635e4c4b0aa0b5b9bb39dde70/tokenizers-0.23.1-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1bf13402aff9bc533c89cb849ec3b412dc3fbeacc9744840e423d7bf3f7dc0e3", size = 3374081, upload-time = "2026-04-27T14:43:01.241Z" }, - { url = "https://files.pythonhosted.org/packages/a2/ef/7735d226f9c7f874a6bee5e3f27fb25ecabdf207d37b8cf45286d0795893/tokenizers-0.23.1-cp310-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:f836ca703b89ae07919a309f9651f7a88fd5a33d5f718ba5ad0870ec0256bad6", size = 3247641, upload-time = "2026-04-27T14:43:03.856Z" }, - { url = "https://files.pythonhosted.org/packages/b9/d9/24827036f6e21297bfffda0768e58eb6096a4f411e932964a01707857931/tokenizers-0.23.1-cp310-abi3-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ae848657742035523fdf261773630cb819a26995fcd3d9ecae0c1daf6e5a4959", size = 3585624, upload-time = "2026-04-27T14:43:10.664Z" }, - { url = "https://files.pythonhosted.org/packages/0c/9a/22f3582b3a4f49358293a5206e25317621ee4526bfe9cdaa0f07a12e770e/tokenizers-0.23.1-cp310-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:53b09e85775d5187941e7bab30e941b4134ab4a7dd8c68e783d231fb7ca27c51", size = 3844062, upload-time = "2026-04-27T14:43:05.643Z" }, - { url = "https://files.pythonhosted.org/packages/7e/65/b8f8814eef95800f20721384136d9a1d22241d50b2874357cb70542c392f/tokenizers-0.23.1-cp310-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:ea5a0ce170074329faaa8ea3f6400ecde604b6678192688533af80980daae71a", size = 3460098, upload-time = "2026-04-27T14:43:08.854Z" }, - { url = "https://files.pythonhosted.org/packages/0d/d5/1353e5f677ec27c2494fb6a6725e82d56c985f53e90ec511369e7e4f02c6/tokenizers-0.23.1-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5075b405006415ea148a992d093699c66eb01952bf59f4d5727089a98bda45a4", size = 3346235, upload-time = "2026-04-27T14:43:12.377Z" }, - { url = "https://files.pythonhosted.org/packages/71/89/39b6b8fc073fb6d413d0147aa333dc7eff7be65639ac9d19930a0b21bf33/tokenizers-0.23.1-cp310-abi3-manylinux_2_31_riscv64.whl", hash = "sha256:56f3a77de629917652f876294dc9fe6bad4a0c43bc229dc72e59bb23a0f4729a", size = 3426398, upload-time = "2026-04-27T14:43:07.264Z" }, - { url = "https://files.pythonhosted.org/packages/0f/80/127c854da64827e5b79264ce524993a90dddcb320e5cd42412c5c02f9e8a/tokenizers-0.23.1-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:9d10a6d957ef01896dc274e890eee27d41bd0e74ef31e60616f0fc311345184e", size = 9823279, upload-time = "2026-04-27T14:43:17.222Z" }, - { url = "https://files.pythonhosted.org/packages/fe/ba/44c2502feb1a058f096ddfb4e0996ef3225a01a388e1a9b094e91689fe93/tokenizers-0.23.1-cp310-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:1974288a609c343774f1b897c8b482c791ab17b75ab5c8c2b1737565c1d82288", size = 9644986, upload-time = "2026-04-27T14:43:19.45Z" }, - { url = "https://files.pythonhosted.org/packages/9e/c1/464019a9fb059870bfe4eebb4ba12208f3042035e258bf5e782906bd3847/tokenizers-0.23.1-cp310-abi3-musllinux_1_2_i686.whl", hash = "sha256:120468fb4c24faf0543c835a4fabafa4deb3f20a035c9b6e83d0b553a97615d4", size = 9976181, upload-time = "2026-04-27T14:43:21.463Z" }, - { url = "https://files.pythonhosted.org/packages/79/94/3ac1432bda31626071e9b6a12709b97ae05131c804b94c8f3ac622c5da32/tokenizers-0.23.1-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:e3d8f40ea6268047de7046906326abed5134f27d4e8447b23763afe5808c8a96", size = 10113853, upload-time = "2026-04-27T14:43:23.617Z" }, - { url = "https://files.pythonhosted.org/packages/6a/dd/631b21433c771b1382535326f0eca80b9c9cee2e64961dd993bc9ac4669e/tokenizers-0.23.1-cp310-abi3-win32.whl", hash = "sha256:93120a930b919416da7cd10a2f606ac9919cc69cacae7980fa2140e277660948", size = 2536263, upload-time = "2026-04-27T14:43:29.888Z" }, - { url = "https://files.pythonhosted.org/packages/97/c9/2553f72aaf65a2797d4229e37fa7fbe38ffbf3e32912d31bdd78b3323e59/tokenizers-0.23.1-cp310-abi3-win_amd64.whl", hash = "sha256:e7bfaf995c1bdbbd21d13539decb6650967013759318627d85daeb7881af16b7", size = 2798223, upload-time = "2026-04-27T14:43:28.51Z" }, - { url = "https://files.pythonhosted.org/packages/cd/2b/2be299bab55fc595e3d38567edb1a87f86e594842968fa9515a07bdcf422/tokenizers-0.23.1-cp310-abi3-win_arm64.whl", hash = "sha256:a26197957d8e4425dfba746315f3c425ea00cfa8367c5fbc4ec73447893dcea9", size = 2664127, upload-time = "2026-04-27T14:43:26.949Z" }, + { url = "https://files.pythonhosted.org/packages/92/97/5dbfabf04c7e348e655e907ed27913e03db0923abb5dfdd120d7b25630e1/tokenizers-0.22.2-cp39-abi3-macosx_10_12_x86_64.whl", hash = "sha256:544dd704ae7238755d790de45ba8da072e9af3eea688f698b137915ae959281c", size = 3100275, upload-time = "2026-01-05T10:41:02.158Z" }, + { url = "https://files.pythonhosted.org/packages/2e/47/174dca0502ef88b28f1c9e06b73ce33500eedfac7a7692108aec220464e7/tokenizers-0.22.2-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:1e418a55456beedca4621dbab65a318981467a2b188e982a23e117f115ce5001", size = 2981472, upload-time = "2026-01-05T10:41:00.276Z" }, + { url = "https://files.pythonhosted.org/packages/d6/84/7990e799f1309a8b87af6b948f31edaa12a3ed22d11b352eaf4f4b2e5753/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2249487018adec45d6e3554c71d46eb39fa8ea67156c640f7513eb26f318cec7", size = 3290736, upload-time = "2026-01-05T10:40:32.165Z" }, + { url = "https://files.pythonhosted.org/packages/78/59/09d0d9ba94dcd5f4f1368d4858d24546b4bdc0231c2354aa31d6199f0399/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:25b85325d0815e86e0bac263506dd114578953b7b53d7de09a6485e4a160a7dd", size = 3168835, upload-time = "2026-01-05T10:40:38.847Z" }, + { url = "https://files.pythonhosted.org/packages/47/50/b3ebb4243e7160bda8d34b731e54dd8ab8b133e50775872e7a434e524c28/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:bfb88f22a209ff7b40a576d5324bf8286b519d7358663db21d6246fb17eea2d5", size = 3521673, upload-time = "2026-01-05T10:40:56.614Z" }, + { url = "https://files.pythonhosted.org/packages/e0/fa/89f4cb9e08df770b57adb96f8cbb7e22695a4cb6c2bd5f0c4f0ebcf33b66/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:1c774b1276f71e1ef716e5486f21e76333464f47bece56bbd554485982a9e03e", size = 3724818, upload-time = "2026-01-05T10:40:44.507Z" }, + { url = "https://files.pythonhosted.org/packages/64/04/ca2363f0bfbe3b3d36e95bf67e56a4c88c8e3362b658e616d1ac185d47f2/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:df6c4265b289083bf710dff49bc51ef252f9d5be33a45ee2bed151114a56207b", size = 3379195, upload-time = "2026-01-05T10:40:51.139Z" }, + { url = "https://files.pythonhosted.org/packages/2e/76/932be4b50ef6ccedf9d3c6639b056a967a86258c6d9200643f01269211ca/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:369cc9fc8cc10cb24143873a0d95438bb8ee257bb80c71989e3ee290e8d72c67", size = 3274982, upload-time = "2026-01-05T10:40:58.331Z" }, + { url = "https://files.pythonhosted.org/packages/1d/28/5f9f5a4cc211b69e89420980e483831bcc29dade307955cc9dc858a40f01/tokenizers-0.22.2-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:29c30b83d8dcd061078b05ae0cb94d3c710555fbb44861139f9f83dcca3dc3e4", size = 9478245, upload-time = "2026-01-05T10:41:04.053Z" }, + { url = "https://files.pythonhosted.org/packages/6c/fb/66e2da4704d6aadebf8cb39f1d6d1957df667ab24cff2326b77cda0dcb85/tokenizers-0.22.2-cp39-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:37ae80a28c1d3265bb1f22464c856bd23c02a05bb211e56d0c5301a435be6c1a", size = 9560069, upload-time = "2026-01-05T10:45:10.673Z" }, + { url = "https://files.pythonhosted.org/packages/16/04/fed398b05caa87ce9b1a1bb5166645e38196081b225059a6edaff6440fac/tokenizers-0.22.2-cp39-abi3-musllinux_1_2_i686.whl", hash = "sha256:791135ee325f2336f498590eb2f11dc5c295232f288e75c99a36c5dbce63088a", size = 9899263, upload-time = "2026-01-05T10:45:12.559Z" }, + { url = "https://files.pythonhosted.org/packages/05/a1/d62dfe7376beaaf1394917e0f8e93ee5f67fea8fcf4107501db35996586b/tokenizers-0.22.2-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:38337540fbbddff8e999d59970f3c6f35a82de10053206a7562f1ea02d046fa5", size = 10033429, upload-time = "2026-01-05T10:45:14.333Z" }, + { url = "https://files.pythonhosted.org/packages/fd/18/a545c4ea42af3df6effd7d13d250ba77a0a86fb20393143bbb9a92e434d4/tokenizers-0.22.2-cp39-abi3-win32.whl", hash = "sha256:a6bf3f88c554a2b653af81f3204491c818ae2ac6fbc09e76ef4773351292bc92", size = 2502363, upload-time = "2026-01-05T10:45:20.593Z" }, + { url = "https://files.pythonhosted.org/packages/65/71/0670843133a43d43070abeb1949abfdef12a86d490bea9cd9e18e37c5ff7/tokenizers-0.22.2-cp39-abi3-win_amd64.whl", hash = "sha256:c9ea31edff2968b44a88f97d784c2f16dc0729b8b143ed004699ebca91f05c48", size = 2747786, upload-time = "2026-01-05T10:45:18.411Z" }, + { url = "https://files.pythonhosted.org/packages/72/f4/0de46cfa12cdcbcd464cc59fde36912af405696f687e53a091fb432f694c/tokenizers-0.22.2-cp39-abi3-win_arm64.whl", hash = "sha256:9ce725d22864a1e965217204946f830c37876eee3b2ba6fc6255e8e903d5fcbc", size = 2612133, upload-time = "2026-01-05T10:45:17.232Z" }, ] [[package]] @@ -3584,6 +3647,26 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/16/e1/3079a9ff9b8e11b846c6ac5c8b5bfb7ff225eee721825310c91b3b50304f/tqdm-4.67.3-py3-none-any.whl", hash = "sha256:ee1e4c0e59148062281c49d80b25b67771a127c85fc9676d3be5f243206826bf", size = 78374, upload-time = "2026-02-03T17:35:50.982Z" }, ] +[[package]] +name = "transformers" +version = "5.12.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "huggingface-hub" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "pyyaml" }, + { name = "regex" }, + { name = "safetensors" }, + { name = "tokenizers" }, + { name = "tqdm" }, + { name = "typer" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/aa/7c/8240f612819718100a9346dc28dea6a11370c3ca9c8c6eabadd3dea4ef29/transformers-5.12.1.tar.gz", hash = "sha256:679ee731c8225347889ad4fb3b2c926a62e9da3b7d284e9d12c791da7272466b", size = 8924054, upload-time = "2026-06-15T17:27:50.604Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/df/56/bbd60dd8668055803bf8ba55a81f9b8a8b31497f620109a9671d26a2076d/transformers-5.12.1-py3-none-any.whl", hash = "sha256:2a5e109d2021265df7098ffbb738295acaf5ad256f12cbc586db2ea4dcbb1a8a", size = 11150587, upload-time = "2026-06-15T17:27:46.679Z" }, +] + [[package]] name = "typer" version = "0.25.1"