@oh-my-pi/pi-agent-core 18.6.0 → 18.6.3

package/src/tokenizer.ts

@@ -5,6 +5,7 @@ import { materializeString, stringifyJson } from "@oh-my-pi/pi-utils";
 import { LRUCache } from "@oh-my-pi/pi-utils/lru";
 import * as snapcompact from "@oh-my-pi/snapcompact";
 import { isEstimateCacheable, messageEstimateVersion } from "./compaction/message-cache";
+import { base64ImageSize, estimateImageContentTokens } from "./image-tokens";
 import type { AgentMessage } from "./types";
 
 const testEnv = Bun.env.NODE_ENV === "test";
@@ -131,12 +132,6 @@ export interface TokenBudgetCheck {
 	exact: boolean;
 }
 
-/**
- * Image content has no tokenizer representation; charge a fixed estimate
- * matching what providers typically bill for inline images.
- */
-const IMAGE_TOKEN_ESTIMATE = 1200;
-
 /**
  * Memoized estimates for one message under this tokenizer's encoding, split by
  * the {@link MessageCountOptions.excludeEncryptedReasoning} option so the two
@@ -154,13 +149,15 @@ interface MessageEstimate {
  * Model-aware local token counter. Immutable: the catalog-resolved encoding
  * is fixed at construction, so a cached count can never straddle two
  * encodings. An `Agent` owns one for its active model (swapping the instance
- * when the model's encoding changes); one-shot flows construct their own for
- * the model that will be billed. Known tokenizer families use exact native
- * counts; unknown models keep the fast byte estimate (or o200k when
- * `PI_TOKENIZER_ACCURATE=1`).
+ * when the model's encoding or snapcompact frame pricing changes); one-shot
+ * flows construct their own for the model that will be billed. Known
+ * tokenizer families use exact native counts; unknown models keep the fast
+ * byte estimate (or o200k when `PI_TOKENIZER_ACCURATE=1`).
  */
 export class Tokenizer {
 	readonly #encoding: natives.Encoding | null;
+	readonly #frameBilling: snapcompact.FrameBilling;
+	readonly frameBillingKey: string;
 
 	/** Exact counts only; byte fallbacks remain mode-dependent and uncached. */
 	readonly #nativeCounts = new LRUCache<string, number>({
@@ -179,8 +176,10 @@ export class Tokenizer {
 	 */
 	#estimates = new WeakMap<AgentMessage, MessageEstimate>();
 
-	constructor(model?: Pick<Model, "tokenizer"> | null) {
+	constructor(model?: (Pick<Model, "tokenizer"> & snapcompact.ShapeTarget) | null) {
 		this.#encoding = tokenizerEncodingForModel(model);
+		this.#frameBilling = snapcompact.frameBilling(model ?? undefined);
+		this.frameBillingKey = snapcompact.frameBillingKey(this.#frameBilling);
 	}
 
 	get encoding(): natives.Encoding | null {
@@ -219,7 +218,8 @@ export class Tokenizer {
 	 * Settled historical messages are counted once and reused until an owner
 	 * (prune/shake/strip-images) calls `invalidateMessageCache`; streaming
 	 * assistants bypass the memo entirely (see the message-cache settle-gate
-	 * invariant). Image blocks charge a fixed per-image estimate.
+	 * invariant). Image blocks charge {@link estimateImageContentTokens}, the
+	 * same dimension-based estimate the native remote-compaction probe applies.
 	 */
 	countMessage(message: AgentMessage, options?: MessageCountOptions): number {
 		const floored = options?.excludeEncryptedReasoning === true;
@@ -273,7 +273,7 @@ export class Tokenizer {
 						if (block.type === "text" && block.text) {
 							fragments.push(block.text);
 						} else if (block.type === "image") {
-							extra += IMAGE_TOKEN_ESTIMATE;
+							extra += estimateImageContentTokens(block);
 						}
 					}
 				}
@@ -323,7 +323,7 @@ export class Tokenizer {
 						if (block.type === "text" && block.text) {
 							fragments.push(block.text);
 						} else if (block.type === "image") {
-							extra += IMAGE_TOKEN_ESTIMATE;
+							extra += estimateImageContentTokens(block);
 						}
 					}
 				}
@@ -336,11 +336,12 @@ export class Tokenizer {
 					if (message.blocks) {
 						for (const block of message.blocks) {
 							if (block.type === "text") fragments.push(block.text);
-							else extra += snapcompact.FRAME_TOKEN_ESTIMATE;
+							else extra += this.#frameTokens(block.data);
 						}
 					} else if (message.images) {
-						// Snapcompact frames render at ≥1568px; providers bill the downscaled cap.
-						extra += message.images.length * snapcompact.FRAME_TOKEN_ESTIMATE;
+						for (const image of message.images) {
+							extra += this.#frameTokens(image.data);
+						}
 					}
 				}
 				break;
@@ -352,4 +353,10 @@ export class Tokenizer {
 		if (fragments.length === 0) return extra;
 		return extra + this.countTokens(fragments);
 	}
+
+	/** One snapcompact frame at the active model's price for its pixel size; unreadable frames cost the ceiling. */
+	#frameTokens(data: string): number {
+		const size = base64ImageSize(data);
+		return size ? snapcompact.frameTokens(this.#frameBilling, size) : snapcompact.FRAME_TOKEN_ESTIMATE;
+	}
 }