@tokz/ai-sdk

Drop-in wrapper for the Vercel AI SDK. It intercepts the tool-to-model boundary: your tools keep returning what they always returned, and the model sees a compressed view of it.

Compression is deterministic and extractive. Hosted mode sends tool output to Tokz without retaining it; inject a local compressor to keep it in process. Set tokz.prepareStep: true for the complete per-step lifecycle. History policies support append-only, idle, and cost-aware compaction for provider prefix caches.

Install

pnpm add @tokz/ai-sdk ai

ai is a peer dependency (>=5). This package builds and typechecks without it; you only need it at runtime, and even then only if you do not pass your own delegate.

Usage

import { generateText } from "@tokz/ai-sdk";
import { openai } from "@ai-sdk/openai";
import { tool } from "ai";
import { z } from "zod";

const listPods = tool({
  description: "List pods in a namespace",
  inputSchema: z.object({ namespace: z.string() }),
  execute: async ({ namespace }) => {
    const res = await fetch(`https://k8s.internal/api/v1/namespaces/${namespace}/pods`);
    return res.json(); // 40 KB of near-identical rows
  },
});

const result = await generateText({
  model: openai("gpt-4o"),
  messages: [{ role: "user", content: "Which pods in prod are unhealthy?" }],
  tools: { listPods },
  tokz: {
    targetRatio: 0.45,
    autoExpand: "on-reference",
    conversationCompaction: true,
  },
});

console.log(result.text);
console.log(result.tokz.originalBytes, "->", result.tokz.renderedBytes);

Everything except tokz is forwarded to ai.generateText untouched, and the result is the SDK's result with one extra field, result.tokz.

The tokz option block

| Option | Default | Effect | | --- | --- | --- | | targetRatio | 0.45 | Share of tool-output bytes to keep. The structural skeleton is never dropped, so the achieved ratio can be higher. | | autoExpand | "never" | "on-reference" recovers elided content the answer refers to and re-prompts once. | | conversationCompaction | false | Re-compresses retained tool results in messages before delegating. | | compactionRatio | 0.6 * targetRatio | Budget used by that compaction. | | store | shared store | Where originals and span maps are retained. Pass one per conversation if a process serves many. | | delegate | ai.generateText | The underlying call. Injectable for tests. |

Recovering the original

The model sees the compressed view while the original stays in your local retention store. Hosted Tokz receives the text during compression but does not retain it:

import { Tokz, expand } from "@tokz/sdk";
import { MapRetentionStore } from "@tokz/ai-sdk";
import type { RetentionRecord } from "@tokz/ai-sdk";

// In real usage this record is populated by generateText's tool wrapping and
// read back via `result.tokz.store.get(toolCallId)`; built by hand here so the
// recovery flow is runnable on its own.
const tokz = new Tokz({ apiKey: process.env.TOKZ_API_KEY! });
const store = new MapRetentionStore();

const originalText = JSON.stringify({
  pods: Array.from({ length: 200 }, (_, i) => ({ name: `web-${i}`, status: "Running" })),
});
const compressed = await tokz.compress(originalText, { targetRatio: 0.3 });

const record: RetentionRecord = {
  toolName: "listPods",
  toolCallId: "call_1",
  originalText,
  result: compressed,
  rendered: compressed.text,
  targetRatio: 0.3,
};
store.set(record);

const stored = store.get("call_1")!;
stored.originalText;                                    // exactly what the tool returned
if (stored.result.method !== "semantic") {
  expand(stored.originalText, stored.result.spanMap);               // verified round-trip
  expand(stored.originalText, stored.result.spanMap, { elision: 0 }); // just the dropped run
}

Compacting history

import { compactHistory, apiCompressor } from "@tokz/ai-sdk";
import type { ModelMessageLike } from "@tokz/ai-sdk";

const compress = apiCompressor({ apiKey: process.env.TOKZ_API_KEY! });
const messages: ModelMessageLike[] = [{ role: "user", content: "Which pods are unhealthy?" }];

const tighter = await compactHistory(messages, 0.2, { compress });

This is a pure local function. The original text and its span map are already in the store, so a smaller view of an old tool result is the same deterministic CPU pass that produced the first view — zero API calls, zero inference. A wrapper that discarded the original would have to ask a model to summarise history, which costs more than the tokens it saves.

Messages are not mutated; changed ones come back as new objects.

Auto-expand is a heuristic

detectReferencedElisions decides that the model "referred to" elided content by vocabulary overlap: it recovers each elided run locally, takes the identifiers in it, and flags the run when the response mentions one at a word boundary. It over-fires on generic words and under-fires when the model paraphrases without naming anything.

Treat a hit as a cheap reason to re-supply context, not as proof. The recovery is exact — expand verifies srcSha256 and returns the dropped bytes verbatim — only the trigger is fuzzy.

import { detectReferencedElisions, expansionMessage, sharedStore } from "@tokz/ai-sdk";
import type { ModelMessageLike } from "@tokz/ai-sdk";

const responseText = "The web-42 pod is crash-looping; see its restart count above.";
const hits = detectReferencedElisions(responseText, sharedStore);

const messages: ModelMessageLike[] = [];
if (hits.length) messages.push(expansionMessage(hits));

What the model actually sees

Wrapped tools set toModelOutput — the AI SDK's documented tool-to-model hook — and return the compressed render from execute. Both, deliberately: the hook is the correct mechanism, and returning the render as well means the original cannot reach a message even on an SDK version that skips the hook. Tools with no server-side execute (client-side or provider-executed) are left alone, as are tools that stream their output as an AsyncIterable.