[Improvement] count_tokens re-encodes the whole transcript every turn (#1736)

[ayaangazali]

Fixes #1736.

What

count_tokens() has no memoization and is called from 22 sites, including the per-turn budget checks (agent.py:check_available_tokens, tokens_checks) and the transforms pipeline — always against the full, growing transcript. Every turn re-encodes the entire prefix it already encoded on the previous turn. That is the engine behind the O(n²) conversation cost in #1735.

Measured on a ~36k-word transcript:

cold  7.500 ms
warm  0.075 ms      ~100x

The saving compounds with conversation length, which is exactly where it is needed.

Fix

A bounded, digest-keyed cache in swarms/utils/litellm_tokenizer.py.

Keyed by SHA-256 of the text, not by the text. This is the part worth reviewing. The obvious implementation is @lru_cache directly on count_tokens, but lru_cache retains its arguments as the key — a few hundred entries of full transcript would pin hundreds of megabytes for a long run, costing more memory than the CPU it saves. Only the 32-byte digest is retained.

The key includes the model, so two models never share an entry.

Only successful counts are cached. The fallback path is left uncached so a transient tokenizer failure is never remembered as an answer, and the empty/whitespace short-circuit still returns 0 before the cache is consulted.

Bounded at 512 entries, LRU-evicted via OrderedDict.move_to_end / popitem(last=False) — no unbounded growth in a long-lived process.

encode is deterministic for a given model, so a hit is exact, not approximate.

Scope

Only the memoization. #1736 also suggests adopting the running-total approach from #1735 so the full history is never re-counted at all; that is a larger change across conversation.py and belongs on its own. The two are complementary — this one helps immediately and does not conflict with that work.

This supersedes #1784, which I closed myself back in August purely to keep my open-PR count down, not on merit.

Test

Appended to tests/utils/test_litellm_wrapper.py — no new file. Nothing owns litellm_tokenizer.py today, and that file is the nearest owner (the litellm utils in swarms/utils/).

It asserts the second call returns the same count faster than the first, that distinct inputs are not conflated, that the cache stops growing at its bound after MAX + 50 distinct inputs, and that keys are model-scoped.

master source + this test:   1 failed
this branch:                 2 passed
black --check --line-length 70, ruff:  clean

Red checks are the repo-wide pre-existing ones.

---

Update — test file dropped. This is a small change to one function, so it ships without a test per the repo's preference for keeping diffs to the fix itself. The verification described above was run directly (reproduction before, same reproduction after); nothing about the fix or the evidence changed, only the absence of a committed test file.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.