From 2b66da71c52ec47d854e02755e1cb1736179e532 Mon Sep 17 00:00:00 2001 From: leetcrypt Date: Fri, 26 Jun 2026 12:27:59 -0700 Subject: [PATCH] feat(ai): calibrate native pruning against real Ollama token counts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 4 (final) of the native-harness working-memory sprint. OllamaProvider .complete_with_tools now returns (text, calls, usage), surfacing the response's real prompt_eval_count/eval_count (free — already in the payload). The native loop EMA-smooths real/estimate into self._tok_ratio (clamped [0.5,3.0]) and prunes against native_token_budget / ratio, so context budgeting tracks the TRUE window instead of the systematic bias of the len//4 char estimate. Providers that omit counts leave the ratio at 1.0, so behaviour is unchanged where unavailable — a free correctness win, no regression. RAM-only, no disk, no new network frames. Scope note: touches only cmd_chat/agent/ + providers — disjoint from the parallel feat(operator) work on this branch, so it merges/reverts independently by path. Sprint: native-harness-working-memory (Phase 4/4) — see also 46e5620, dc6317f Co-Authored-By: Claude Opus 4.6 --- cmd_chat/agent/bridge.py | 30 +++++++++++++++++++++++++++--- cmd_chat/agent/providers.py | 20 ++++++++++++++------ 2 files changed, 41 insertions(+), 9 deletions(-) diff --git a/cmd_chat/agent/bridge.py b/cmd_chat/agent/bridge.py index 84c309a..5fd50e3 100644 --- a/cmd_chat/agent/bridge.py +++ b/cmd_chat/agent/bridge.py @@ -328,6 +328,12 @@ class AgentBridge(Client): # probe. (cwd, shell) once learned; None until the first native run. Dies # with the agent — no disk, consistent with the rest of the agent's memory. self._sbx_known: tuple[str, str] | None = None + # Calibration of the char-based token estimate against Ollama's REAL + # prompt_eval_count (Phase 4). real ≈ ratio × estimate; the native loop + # budgets pruning against `native_token_budget / ratio` so it tracks the + # true context window instead of the systematic bias of `len//4`. EMA- + # smoothed and clamped; 1.0 until the first turn yields a real count. + self._tok_ratio: float = 1.0 @staticmethod def _est_tokens(text: str) -> int: @@ -1153,8 +1159,11 @@ class AgentBridge(Client): # Keep the running context under budget: drop the oldest middle turns # (task + freshest turns pinned) so the goal never falls out of a weak # model's effective window mid-task. Logged, never silent. - messages, pruned_n, digested_n = self._prune_native_messages( - messages, self.native_token_budget) + # Budget against TRUE tokens (Phase 4): real ≈ ratio × estimate, so prune + # to keep the char-estimate under `budget / ratio`. Clamp the divisor so a + # noisy calibration sample can neither collapse nor balloon the window. + eff_budget = int(self.native_token_budget / min(max(self._tok_ratio, 0.5), 3.0)) + messages, pruned_n, digested_n = self._prune_native_messages(messages, eff_budget) if pruned_n or digested_n: detail = [] if digested_n: @@ -1177,7 +1186,7 @@ class AgentBridge(Client): turn_system += "\n\n" + ws_block await self._send_typing(ws, True) try: - text, calls = await asyncio.to_thread(cwt, turn_system, messages, NATIVE_TOOLS) + text, calls, usage = await asyncio.to_thread(cwt, turn_system, messages, NATIVE_TOOLS) except ToolsUnsupported: await self._send_typing(ws, False) await self._send_chat(ws, f"{asker}: (model has no tool support — using simple harness)") @@ -1189,6 +1198,21 @@ class AgentBridge(Client): return await self._send_typing(ws, False) + # Calibrate the char estimator against Ollama's real prompt_eval_count + # (the prompt it just processed = system + tools + messages). EMA-smoothed + # and clamped so pruning tracks the true window without chasing jitter. + # Providers that omit counts leave the ratio — and thus behaviour — + # unchanged, so this is a free correctness win where available. + pt = usage.get("prompt_eval_count") if usage else None + if pt: + est_prompt = (self._est_tokens(turn_system) + + self._est_tokens(json.dumps(NATIVE_TOOLS)) + + sum(self._est_tokens(str(m.get("content") or "")) + for m in messages)) + if est_prompt > 0: + sample = pt / est_prompt + self._tok_ratio = min(max(0.7 * self._tok_ratio + 0.3 * sample, 0.5), 3.0) + if not calls and not self._is_done_signal(text): # The weak models' dominant failure is narrating the next command in a # ```bash fence instead of calling run_shell. Recover that as a real diff --git a/cmd_chat/agent/providers.py b/cmd_chat/agent/providers.py index b043c31..eb83fa2 100644 --- a/cmd_chat/agent/providers.py +++ b/cmd_chat/agent/providers.py @@ -128,13 +128,16 @@ class OllamaProvider: def complete_with_tools( self, system: str, messages: list[dict], tools: list[dict] - ) -> tuple[str, list[dict]]: + ) -> tuple[str, list[dict], dict]: """One non-streaming ``/api/chat`` turn carrying a ``tools`` schema. Used by the native harness loop. ``messages`` are raw Ollama wire dicts (so the caller can round-trip assistant ``tool_calls`` and ``tool`` results across - turns); ``system`` is prepended. Returns ``(text, tool_calls)`` where each - call is ``{"name": str, "arguments": dict}``. Raises ``ToolsUnsupported`` if - the model can't do function calling so the bridge can fall back to simple.""" + turns); ``system`` is prepended. Returns ``(text, tool_calls, usage)`` where + each call is ``{"name": str, "arguments": dict}`` and ``usage`` carries + Ollama's real token counts (``prompt_eval_count`` / ``eval_count``) so the + caller can budget context against TRUE tokens instead of a char estimate + (``{}`` if the server omits them). Raises ``ToolsUnsupported`` if the model + can't do function calling so the bridge can fall back to simple.""" # Greedy decode (temperature 0) for the tool loop: at Ollama's default 0.8 a # weak model "creatively" narrates the next step in prose or fabricates file # content instead of emitting a deterministic structured call. The nudge loop @@ -161,7 +164,12 @@ class OllamaProvider: raise ToolsUnsupported(detail or f"{self.model} does not support tools") self._raise_for_status(r) self._tools_ok = True - msg = r.json().get("message", {}) or {} + data = r.json() + msg = data.get("message", {}) or {} + # Real token counts straight from Ollama — exact, free (already in the + # response), and used to calibrate the native loop's char-based estimate. + usage = {k: data[k] for k in ("prompt_eval_count", "eval_count") + if isinstance(data.get(k), int)} text = (msg.get("content") or "").strip() calls: list[dict] = [] for tc in msg.get("tool_calls") or []: @@ -188,7 +196,7 @@ class OllamaProvider: recovered_text, recovered = self._extract_text_tool_calls(text, valid) if recovered: text, calls = recovered_text, recovered - return text, calls + return text, calls, usage # Wrapper tags a weak model wraps a leaked call (or its prose) in; stripped # from the chat-facing text once the JSON inside is recovered.