diff --git a/broker.lua b/broker.lua index a3e9e96..2527200 100644 --- a/broker.lua +++ b/broker.lua @@ -229,4 +229,58 @@ function M.chat(model_cfg, messages, opts) return table.concat(parts), captured_usage end +-- ---------------------------------------------------------------- token_count (Phase 8) +-- Returns an accurate token count by hitting /tokenize when +-- the endpoint supports it; falls back to the Phase 0 ยง8 char/4 +-- heuristic otherwise. Per-endpoint capability cache (session-local; +-- key per R6 is endpoint-only since B1 confirms /tokenize ignores the +-- model field on the observed broker). +-- +-- Never errors. Returns a non-negative integer. +-- 2s timeout per call so a misbehaving endpoint can't stall the +-- caller; first miss caches as unsupported for the session. +local _tokenize_capable = {} -- [endpoint] = true | false (nil = unprobed) + +function M.token_count(model_cfg, text) + text = text or "" + if text == "" then return 0 end + if not (model_cfg and model_cfg.endpoint) then + return math.floor(#text / 4) + end + local ep = model_cfg.endpoint + local cap = _tokenize_capable[ep] + if cap == false then + return math.floor(#text / 4) + end + local url = ep:gsub("/+$", "") .. "/tokenize" + local body = json.encode({ content = text, model = model_cfg.model }) + local out, status = curl.post(url, body, + { "Content-Type: application/json" }, + 2000) -- 2s timeout per R5 risk row + if not (status == 200 and out) then + _tokenize_capable[ep] = false + return math.floor(#text / 4) + end + local doc = json.decode(out) + local toks = doc and doc.tokens + if type(toks) ~= "table" then + _tokenize_capable[ep] = false + return math.floor(#text / 4) + end + _tokenize_capable[ep] = true + return #toks +end + +-- Introspection: nil if endpoint un-probed; true/false for the cached +-- capability. Used by tests and future :tokenize debug meta. +function M.tokenize_supported(model_cfg) + if not (model_cfg and model_cfg.endpoint) then return nil end + return _tokenize_capable[model_cfg.endpoint] +end + +-- Test hook: reset the cache between LuaJIT-VM-shared test runs. +function M._reset_tokenize_cache() + _tokenize_capable = {} +end + return M