From 88e08cdbe5de8b6e2b45a4b2e3e47c7d9c171726 Mon Sep 17 00:00:00 2001 From: Vignesh Date: Wed, 1 Apr 2026 19:36:26 +0530 Subject: [PATCH 1/2] chore(agents): add js-tiktoken for model-aware token counting --- bun.lock | 5 +++++ packages/agents/package.json | 3 ++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/bun.lock b/bun.lock index c365814..3f6545f 100644 --- a/bun.lock +++ b/bun.lock @@ -50,6 +50,7 @@ "version": "0.1.0", "dependencies": { "@stackforge/shared": "workspace:*", + "js-tiktoken": "^1.0.21", }, "devDependencies": { "@types/bun": "^1.3.11", @@ -336,6 +337,8 @@ "array-flatten": ["array-flatten@1.1.1", "", {}, "sha512-PCVAQswWemu6UdxsDFFX/+gVeYqKAod3D3UVm91jHwynguOwAvYPhx8nNlM++NqRcK6CxxpUafjmhIdKiHibqg=="], + "base64-js": ["base64-js@1.5.1", "", {}, "sha512-AKpaYlHn8t4SVbOHCy+b5+KKgvR4vrsD8vbvrbiQJps7fKDTkjkDry6ji0rUJjC0kzbNePLwzxq8iypo41qeWA=="], + "baseline-browser-mapping": ["baseline-browser-mapping@2.10.13", "", { "bin": { "baseline-browser-mapping": "dist/cli.cjs" } }, "sha512-BL2sTuHOdy0YT1lYieUxTw/QMtPBC3pmlJC6xk8BBYVv6vcw3SGdKemQ+Xsx9ik2F/lYDO9tqsFQH1r9PFuHKw=="], "body-parser": ["body-parser@1.20.4", "", { "dependencies": { "bytes": "~3.1.2", "content-type": "~1.0.5", "debug": "2.6.9", "depd": "2.0.0", "destroy": "~1.2.0", "http-errors": "~2.0.1", "iconv-lite": "~0.4.24", "on-finished": "~2.4.1", "qs": "~6.14.0", "raw-body": "~2.5.3", "type-is": "~1.6.18", "unpipe": "~1.0.0" } }, "sha512-ZTgYYLMOXY9qKU/57FAo8F+HA2dGX7bqGc71txDRC1rS4frdFI5R7NhluHxH6M0YItAP0sHB4uqAOcYKxO6uGA=="], @@ -438,6 +441,8 @@ "jiti": ["jiti@2.6.1", "", { "bin": { "jiti": "lib/jiti-cli.mjs" } }, "sha512-ekilCSN1jwRvIbgeg/57YFh8qQDNbwDb9xT/qu2DAHbFFZUicIl4ygVaAvzveMhMVr3LnpSKTNnwt8PoOfmKhQ=="], + "js-tiktoken": ["js-tiktoken@1.0.21", "", { "dependencies": { "base64-js": "^1.5.1" } }, "sha512-biOj/6M5qdgx5TKjDnFT1ymSpM5tbd3ylwDtrQvFQSu0Z7bBYko2dF+W/aUkXUPuk6IVpRxk/3Q2sHOzGlS36g=="], + "js-tokens": ["js-tokens@4.0.0", "", {}, "sha512-RdJUflcE3cUzKiMqQgsCu06FPu9UdIJO0beYbPhHN4k6apgJtifcoCtT9bcxOpYBtpD2kCM6Sbzg4CausW/PKQ=="], "jsesc": ["jsesc@3.1.0", "", { "bin": { "jsesc": "bin/jsesc" } }, "sha512-/sM3dO2FOzXjKQhJuo0Q173wf2KOo8t4I8vHy6lF9poUp7bKT0/NHE8fPX23PwfhnykfqnC2xRxOnVw5XuGIaA=="], diff --git a/packages/agents/package.json b/packages/agents/package.json index 4397514..86301f6 100644 --- a/packages/agents/package.json +++ b/packages/agents/package.json @@ -18,7 +18,8 @@ "test": "bun test" }, "dependencies": { - "@stackforge/shared": "workspace:*" + "@stackforge/shared": "workspace:*", + "js-tiktoken": "^1.0.21" }, "devDependencies": { "@types/bun": "^1.3.11", From 527aef19db7eb7d666dd79bb835196eb26a537be Mon Sep 17 00:00:00 2001 From: Vignesh Date: Wed, 1 Apr 2026 19:36:50 +0530 Subject: [PATCH 2/2] feat(agents): maximize prompt compression with tokenizer-aware budgeting --- packages/agents/src/agents/prompts/index.ts | 6 +- .../agents/src/optimizer/token.optimizer.ts | 317 +++++++++++++----- packages/agents/test/token-optimizer.test.ts | 41 ++- 3 files changed, 284 insertions(+), 80 deletions(-) diff --git a/packages/agents/src/agents/prompts/index.ts b/packages/agents/src/agents/prompts/index.ts index 58c5e57..2ebb19f 100644 --- a/packages/agents/src/agents/prompts/index.ts +++ b/packages/agents/src/agents/prompts/index.ts @@ -59,8 +59,12 @@ const CONTEXT_KEYS: Record = { ], }; +export function extractAgentContext(agentName: AgentName, input: unknown): Record { + return pick(input, CONTEXT_KEYS[agentName]); +} + export function buildAgentPrompt(agentName: AgentName, input: unknown): AgentPrompt { - const context = pick(input, CONTEXT_KEYS[agentName]); + const context = extractAgentContext(agentName, input); return { systemPrompt: diff --git a/packages/agents/src/optimizer/token.optimizer.ts b/packages/agents/src/optimizer/token.optimizer.ts index eee5b09..d8cc5c5 100644 --- a/packages/agents/src/optimizer/token.optimizer.ts +++ b/packages/agents/src/optimizer/token.optimizer.ts @@ -1,6 +1,7 @@ import type { AgentName } from "@stackforge/shared"; +import { encodingForModel, getEncoding, type TiktokenModel, type Tiktoken } from "js-tiktoken"; import { AGENT_CONFIGS } from "../config/agent.configs.js"; -import { buildAgentPrompt } from "../agents/prompts/index.js"; +import { buildAgentPrompt, extractAgentContext } from "../agents/prompts/index.js"; export type OptimizedAgentPayload = { optimizedInput: unknown; @@ -14,8 +15,56 @@ export type OptimizedAgentPayload = { compressionPasses: number; }; -function estimateTokens(text: string): number { - return Math.ceil(text.length / 4); +const tokenizerCache = new Map(); + +function normalizeModel(model: string): string { + const providerStripped = model.includes("/") ? model.split("/").at(-1) ?? model : model; + return providerStripped.split(":")[0]?.trim() ?? providerStripped; +} + +function getTokenizer(model: string): Tiktoken { + const normalized = normalizeModel(model); + const cached = tokenizerCache.get(normalized); + if (cached !== undefined) { + return cached; + } + + let tokenizer: Tiktoken; + try { + tokenizer = encodingForModel(normalized as TiktokenModel); + } catch { + tokenizer = getEncoding("o200k_base"); + } + + tokenizerCache.set(normalized, tokenizer); + return tokenizer; +} + +function estimateTextTokens(text: string, model: string): number { + return getTokenizer(model).encode(text).length; +} + +function estimatePromptTokens( + systemPrompt: string, + userPrompt: string, + model: string, + inputLimit?: number, +): number { + const roughEstimate = Math.ceil((systemPrompt.length + userPrompt.length) / 3.7) + 20; + + if (inputLimit !== undefined && roughEstimate > inputLimit * 1.5) { + return roughEstimate; + } + + const systemTokens = estimateTextTokens(systemPrompt, model); + const userTokens = estimateTextTokens(userPrompt, model); + + // Account for chat message wrappers and small protocol overhead. + return systemTokens + userTokens + 20; +} + +function normalizeWhitespace(value: string): string { + return value.replace(/\s+/g, " ").trim(); } function trimPromptValue(value: unknown, maxChars: number): unknown { @@ -23,137 +72,251 @@ function trimPromptValue(value: unknown, maxChars: number): unknown { return value; } - if (value.length <= maxChars) { - return value; + const normalized = normalizeWhitespace(value); + + if (normalized.length <= maxChars) { + return normalized; + } + + if (maxChars <= 48) { + return normalized.slice(0, maxChars); } - return value.slice(0, maxChars); + const head = Math.floor(maxChars * 0.72); + const tail = Math.max(12, maxChars - head - 5); + return `${normalized.slice(0, head)} ... ${normalized.slice(-tail)}`; } type CompressionPlan = { - promptRatio: number; - genericStringRatio: number; + promptChars: number; + genericStringChars: number; arrayLimit: number; + objectKeyLimit: number; + maxDepth: number; }; function buildCompressionPlan( level: "low" | "medium" | "high", pass: number, + maxInputTokens: number, ): CompressionPlan { + const maxChars = maxInputTokens * 4; + const presets: Record<"low" | "medium" | "high", CompressionPlan> = { - low: { promptRatio: 0.9, genericStringRatio: 0.8, arrayLimit: 24 }, - medium: { promptRatio: 0.75, genericStringRatio: 0.6, arrayLimit: 16 }, - high: { promptRatio: 0.6, genericStringRatio: 0.4, arrayLimit: 10 }, + low: { + promptChars: Math.floor(maxChars * 0.65), + genericStringChars: Math.floor(maxChars * 0.24), + arrayLimit: 28, + objectKeyLimit: 18, + maxDepth: 6, + }, + medium: { + promptChars: Math.floor(maxChars * 0.52), + genericStringChars: Math.floor(maxChars * 0.18), + arrayLimit: 18, + objectKeyLimit: 14, + maxDepth: 5, + }, + high: { + promptChars: Math.floor(maxChars * 0.42), + genericStringChars: Math.floor(maxChars * 0.13), + arrayLimit: 12, + objectKeyLimit: 10, + maxDepth: 4, + }, }; const base = presets[level]; - const adjustment = pass * 0.12; + const scalar = Math.max(0.25, 1 - pass * 0.2); + return { - promptRatio: Math.max(0.2, base.promptRatio - adjustment), - genericStringRatio: Math.max(0.15, base.genericStringRatio - adjustment), + promptChars: Math.max(80, Math.floor(base.promptChars * scalar)), + genericStringChars: Math.max(48, Math.floor(base.genericStringChars * scalar)), arrayLimit: Math.max(3, base.arrayLimit - pass * 4), + objectKeyLimit: Math.max(4, base.objectKeyLimit - pass * 2), + maxDepth: Math.max(2, base.maxDepth - pass), }; } -function clampJsonSize(serialized: string, maxInputTokens: number): string { - const maxChars = maxInputTokens * 4; - return serialized.length <= maxChars ? serialized : serialized.slice(0, maxChars); +const KEY_PRIORITY: Record = { + prompt: 100, + projectName: 98, + stack: 96, + entities: 95, + relationships: 95, + routePlan: 94, + frontendPages: 93, + infraPlan: 92, + generatedFilesPlan: 92, + reviewerNotes: 92, + name: 90, + tableName: 88, + fields: 88, + method: 87, + path: 87, + type: 86, + description: 85, + auth: 84, + responseType: 84, +}; + +function keyScore(key: string): number { + if (key in KEY_PRIORITY) { + return KEY_PRIORITY[key] ?? 0; + } + + if (key.endsWith("Id") || key === "id") { + return 70; + } + + return 30; } -function compressLargeCollections( - value: Record, - arrayLimit: number, -): Record { - const keys = [ - "entities", - "relationships", - "routePlan", - "frontendPages", - "generatedFilesPlan", - "folderStructure", - "reviewerNotes", - "ci", - "deployment", - "envVars", - ]; - - const copy = { ...value }; - for (const key of keys) { - const current = copy[key]; - if (Array.isArray(current) && current.length > arrayLimit) { - copy[key] = current.slice(0, arrayLimit); - } +function sampleArray(values: T[], limit: number): T[] { + if (values.length <= limit) { + return values; } - return copy; + if (limit <= 1) { + return [values[0]!]; + } + + const sampled: T[] = []; + const step = (values.length - 1) / (limit - 1); + for (let index = 0; index < limit; index++) { + sampled.push(values[Math.round(index * step)]!); + } + + return sampled; } -function compressInput(input: unknown, maxInputTokens: number, plan: CompressionPlan): unknown { - const maxChars = maxInputTokens * 4; +function compressRecursive( + value: unknown, + plan: CompressionPlan, + key: string, + depth: number, +): unknown { + if (value === null || value === undefined) { + return value; + } - if (input === null || typeof input !== "object") { - return trimPromptValue(input, maxChars); + if (typeof value === "string") { + return trimPromptValue( + value, + key === "prompt" ? plan.promptChars : plan.genericStringChars, + ); } - let shallow = { ...(input as Record) }; - shallow = compressLargeCollections(shallow, plan.arrayLimit); + if (typeof value === "number" || typeof value === "boolean") { + return value; + } - if ("prompt" in shallow) { - shallow["prompt"] = trimPromptValue(shallow["prompt"], Math.floor(maxChars * plan.promptRatio)); + if (Array.isArray(value)) { + const sampled = sampleArray(value, plan.arrayLimit); + return sampled.map((item) => compressRecursive(item, plan, key, depth + 1)); } - for (const [key, value] of Object.entries(shallow)) { - if (key !== "prompt") { - shallow[key] = trimPromptValue(value, Math.floor(maxChars * plan.genericStringRatio)); + if (typeof value === "object") { + if (depth >= plan.maxDepth) { + return "[truncated]"; + } + + const entries = Object.entries(value as Record) + .sort((a, b) => keyScore(b[0]) - keyScore(a[0])) + .slice(0, plan.objectKeyLimit); + + const out: Record = {}; + for (const [entryKey, entryValue] of entries) { + out[entryKey] = compressRecursive(entryValue, plan, entryKey, depth + 1); } + + return out; } - let serialized = JSON.stringify(shallow); - if (serialized.length <= maxChars) { - return shallow; + return value; +} + +function compressInput(input: unknown, maxInputTokens: number, plan: CompressionPlan): unknown { + return compressRecursive(input, plan, "", 0); +} + +function buildMinimalContext(context: unknown, maxInputTokens: number): unknown { + if (context === null || typeof context !== "object") { + return trimPromptValue(context, maxInputTokens * 4); } - if ("prompt" in shallow) { - const prompt = shallow["prompt"]; - if (typeof prompt === "string") { - shallow["prompt"] = trimPromptValue(prompt, Math.floor(maxChars * Math.max(0.2, plan.promptRatio - 0.2))); - } + const source = context as Record; + const minimal: Record = {}; + + if ("projectName" in source) { + minimal["projectName"] = trimPromptValue(source["projectName"], 80); + } + + if ("stack" in source) { + minimal["stack"] = source["stack"]; } - serialized = JSON.stringify(shallow); - if (serialized.length <= maxChars) { - return shallow; + if ("prompt" in source) { + minimal["prompt"] = trimPromptValue(source["prompt"], Math.floor(maxInputTokens * 3.2)); } - return { prompt: clampJsonSize(serialized, maxInputTokens) }; + return minimal; } export function optimizeAgentPayload(agentName: AgentName, input: unknown): OptimizedAgentPayload { const config = AGENT_CONFIGS[agentName]; const attempts = Math.max(1, config.budgetOverflowRetries + 1); + const effectiveInputLimit = Math.max(64, config.maxInputTokens - 8); + const context = extractAgentContext(agentName, input); - let selectedInput: unknown = input; - let selectedPrompt = JSON.stringify(input); - let estimatedInputTokens = estimateTokens(selectedPrompt); + let selectedInput: unknown = context; + let selectedPrompt = buildAgentPrompt(agentName, selectedInput); + let estimatedInputTokens = estimatePromptTokens( + selectedPrompt.systemPrompt, + selectedPrompt.userPrompt, + config.model, + config.maxInputTokens, + ); let selectedPass = 1; for (let pass = 0; pass < attempts; pass++) { - const plan = buildCompressionPlan(config.compressionLevel, pass); - const compressedInput = compressInput(input, config.maxInputTokens, plan); - const compactInput = clampJsonSize(JSON.stringify(compressedInput), config.maxInputTokens); - const tokenEstimate = estimateTokens(compactInput); + const plan = buildCompressionPlan(config.compressionLevel, pass, config.maxInputTokens); + const compressedInput = compressInput(context, config.maxInputTokens, plan); + const prompt = buildAgentPrompt(agentName, compressedInput); + const tokenEstimate = estimatePromptTokens( + prompt.systemPrompt, + prompt.userPrompt, + config.model, + config.maxInputTokens, + ); selectedInput = compressedInput; - selectedPrompt = compactInput; + selectedPrompt = prompt; estimatedInputTokens = tokenEstimate; selectedPass = pass + 1; - if (tokenEstimate <= config.maxInputTokens) { + if (tokenEstimate <= effectiveInputLimit) { break; } } - if (estimatedInputTokens > config.maxInputTokens) { + if (estimatedInputTokens > effectiveInputLimit) { + const minimal = buildMinimalContext(context, config.maxInputTokens); + const prompt = buildAgentPrompt(agentName, minimal); + const tokenEstimate = estimatePromptTokens( + prompt.systemPrompt, + prompt.userPrompt, + config.model, + config.maxInputTokens, + ); + + selectedInput = minimal; + selectedPrompt = prompt; + estimatedInputTokens = tokenEstimate; + selectedPass = attempts + 1; + } + + if (estimatedInputTokens > effectiveInputLimit) { throw new Error( `Token budget exceeded for agent '${agentName}' after ${attempts} compression attempts`, ); @@ -168,12 +331,10 @@ export function optimizeAgentPayload(agentName: AgentName, input: unknown): Opti const cappedOutputTokens = Math.min(config.maxOutputTokens, remainingBudget); - const prompt = buildAgentPrompt(agentName, selectedInput); - return { optimizedInput: selectedInput, - systemPrompt: prompt.systemPrompt, - userPrompt: prompt.userPrompt, + systemPrompt: selectedPrompt.systemPrompt, + userPrompt: selectedPrompt.userPrompt, model: config.model, maxInputTokens: config.maxInputTokens, maxOutputTokens: cappedOutputTokens, diff --git a/packages/agents/test/token-optimizer.test.ts b/packages/agents/test/token-optimizer.test.ts index 045eb6a..316f7b5 100644 --- a/packages/agents/test/token-optimizer.test.ts +++ b/packages/agents/test/token-optimizer.test.ts @@ -1,5 +1,6 @@ /// import { describe, it, expect } from "bun:test"; +import { encodingForModel } from "js-tiktoken"; import { optimizeAgentPayload } from "../src/optimizer/token.optimizer.js"; import { AGENT_CONFIGS } from "../src/config/agent.configs.js"; @@ -15,7 +16,45 @@ describe("Token optimizer", () => { expect(result.maxInputTokens).toBe(900); expect(result.estimatedInputTokens).toBeLessThanOrEqual(result.maxInputTokens); expect(result.compressionPasses).toBeGreaterThanOrEqual(1); - expect(result.userPrompt.length).toBeLessThanOrEqual(result.maxInputTokens * 4); + + const tokenizer = encodingForModel("gpt-4o-mini"); + const measuredTokens = + tokenizer.encode(result.systemPrompt).length + + tokenizer.encode(result.userPrompt).length + + 20; + expect(measuredTokens).toBeLessThanOrEqual(result.maxInputTokens); + }); + + it("recursively compresses nested arrays and objects", () => { + const hugeEntities = Array.from({ length: 28 }, (_, entityIndex) => ({ + name: `Entity${entityIndex}`, + tableName: `entity_${entityIndex}`, + fields: Array.from({ length: 24 }, (_, fieldIndex) => ({ + name: `field_${fieldIndex}`, + type: "varchar(255)", + description: "x".repeat(180), + nullable: fieldIndex % 2 === 0, + })), + description: "This is a large entity description that repeats. ".repeat(24), + })); + + const result = optimizeAgentPayload("reviewer", { + prompt: "Review this architecture for consistency and security.", + projectName: "compress-test", + stack: { frontend: "react", backend: "express", database: "postgres" }, + entities: hugeEntities, + relationships: [], + routePlan: [], + frontendPages: [], + infraPlan: { ci: ["test"], docker: true, deployment: ["railway"], envVars: ["DATABASE_URL"] }, + generatedFilesPlan: [], + }); + + const optimized = result.optimizedInput as { entities?: Array<{ fields?: unknown[] }> }; + expect(Array.isArray(optimized.entities)).toBe(true); + expect((optimized.entities ?? []).length).toBeLessThanOrEqual(12); + expect(((optimized.entities ?? [])[0]?.fields ?? []).length).toBeLessThanOrEqual(12); + expect(result.estimatedInputTokens).toBeLessThanOrEqual(result.maxInputTokens); }); it("fails fast when remaining budget cannot satisfy minimum output tokens", () => {