the-token-company 0.2.5 → 0.3.1

package/dist/anthropic.js

@@ -1,6 +1,67 @@
 import { TheTokenCompany } from "./client.js";
 import { StatsTTC, compressAnthropicMessages, resolveAggressiveness } from "./compress.js";
 import { CompressionStats } from "./types.js";
+const TTC_SEARCH_TOOL = {
+    name: "ttc_web_search",
+    description: "Search the web for current information. Use this when you need up-to-date facts, prices, news, or any information that may have changed after your training cutoff.",
+    input_schema: {
+        type: "object",
+        properties: {
+            query: {
+                type: "string",
+                description: "The search query"
+            }
+        },
+        required: ["query"]
+    }
+};
+function injectSearchTool(params) {
+    let tools = params.tools ? [...params.tools] : [];
+    tools = tools.filter((t) => !String(t.type ?? "").startsWith("web_search_"));
+    if (!tools.some((t) => t.name === "ttc_web_search")) {
+        tools.push(TTC_SEARCH_TOOL);
+    }
+    return { ...params, tools };
+}
+function hasSearchToolUse(response) {
+    if (response.stop_reason !== "tool_use")
+        return false;
+    return response.content?.some((b) => b.type === "tool_use" && b.name === "ttc_web_search");
+}
+function formatSearchResults(results) {
+    return results.map(r => `Source: ${r.title}\nURL: ${r.url}\n${r.content}`).join("\n\n");
+}
+async function handleSearchLoop(response, params, originalCreate, ttcClient, rest) {
+    while (hasSearchToolUse(response)) {
+        const messages = [...(params.messages || [])];
+        // Add assistant response
+        const assistantContent = response.content.map((b) => {
+            if (typeof b.toJSON === 'function')
+                return b.toJSON();
+            return b;
+        });
+        messages.push({ role: "assistant", content: assistantContent });
+        // Build tool results
+        const toolResults = [];
+        for (const block of response.content) {
+            if (block.type === "tool_use" && block.name === "ttc_web_search") {
+                const query = block.input?.query || "";
+                const searchResult = await ttcClient.search(query);
+                toolResults.push({
+                    type: "tool_result",
+                    tool_use_id: block.id,
+                    content: formatSearchResults(searchResult.results),
+                });
+            }
+        }
+        if (toolResults.length > 0) {
+            messages.push({ role: "user", content: toolResults });
+        }
+        const newParams = { ...params, messages };
+        response = await originalCreate(newParams, ...rest);
+    }
+    return response;
+}
 /**
  * Wrap an Anthropic client to auto-compress non-assistant messages.
  *
@@ -24,25 +85,43 @@ import { CompressionStats } from "./types.js";
  */
 export function withCompression(client, options) {
     const stats = new CompressionStats();
-    const compressor = new StatsTTC(new TheTokenCompany({ apiKey: options.compressionApiKey, appId: options.appId, fetch: options.fetch }), stats);
+    const ttcClient = new TheTokenCompany({ apiKey: options.compressionApiKey, baseUrl: options.baseUrl, appId: options.appId, fetch: options.fetch });
+    const compressor = new StatsTTC(ttcClient, stats);
     const model = options.model ?? "bear-2";
     const roleAggr = resolveAggressiveness(options.aggressiveness ?? 0.2);
+    if (options.compressAssistant && !("assistant" in roleAggr)) {
+        roleAggr["assistant"] = roleAggr["user"] ?? 0.2;
+    }
     const systemAggr = roleAggr["system"];
+    const stripServerToolResults = options.stripServerToolResults ?? false;
+    const webSearch = options.webSearch ?? false;
     const originalCreate = client.messages.create.bind(client.messages);
     client.messages.create = async function (params, ...rest) {
         stats._startTurn();
         if (params?.messages) {
             params = {
                 ...params,
-                messages: await compressAnthropicMessages(compressor, params.messages, model, roleAggr),
+                messages: await compressAnthropicMessages(compressor, params.messages, model, roleAggr, {
+                    stripServerToolResults,
+                    skipToolName: webSearch ? "ttc_web_search" : undefined,
+                }),
             };
         }
         if (systemAggr != null && typeof params?.system === "string" && params.system.trim()) {
             const result = await compressor.compress(params.system, { model, aggressiveness: systemAggr });
             params = { ...params, system: result.output };
         }
+        // Inject search tool
+        if (webSearch) {
+            params = injectSearchTool(params);
+        }
         stats._endTurn();
-        return originalCreate(params, ...rest);
+        let response = await originalCreate(params, ...rest);
+        // Handle search tool loop
+        if (webSearch) {
+            response = await handleSearchLoop(response, params, originalCreate, ttcClient, rest);
+        }
+        return response;
     };
     client.compression = stats;
     return client;

package/dist/client.d.ts

@@ -1,4 +1,4 @@
-import type { CompressResult, TheTokenCompanyOptions } from "./types.js";
+import type { CompressResult, SearchRequestOptions, SearchResult, TheTokenCompanyOptions } from "./types.js";
 export declare const BEAR_1 = "bear-1";
 export declare const BEAR_1_1 = "bear-1.1";
 export declare const BEAR_1_2 = "bear-1.2";
@@ -16,5 +16,6 @@ export declare class TheTokenCompany {
         aggressiveness?: number;
         appId?: string;
     }): Promise<CompressResult>;
+    search(query: string, options?: SearchRequestOptions): Promise<SearchResult>;
     private parseError;
 }

package/dist/client.js

@@ -69,6 +69,50 @@ export class TheTokenCompany {
             compressionRatio: data.output_tokens === 0 ? 0 : data.original_input_tokens / data.output_tokens,
         };
     }
+    async search(query, options = {}) {
+        const { maxResults = 5, searchDepth = "basic", includeRawContent = false, model = "bear-2", aggressiveness = 0.3, } = options;
+        const resolvedAppId = options.appId ?? this.appId;
+        const payload = {
+            query,
+            max_results: maxResults,
+            search_depth: searchDepth,
+            include_raw_content: includeRawContent,
+            model,
+            compression_settings: { aggressiveness },
+            ...(resolvedAppId !== undefined && { app_id: resolvedAppId }),
+        };
+        const jsonBody = JSON.stringify(payload);
+        const headers = {
+            Authorization: `Bearer ${this.apiKey}`,
+            "Content-Type": "application/json",
+        };
+        let body;
+        if (this.gzip) {
+            headers["Content-Encoding"] = "gzip";
+            body = new Uint8Array(await gzipAsync(jsonBody));
+        }
+        else {
+            body = jsonBody;
+        }
+        const response = await this.fetchFn(`${this.baseUrl}/v1/search`, {
+            method: "POST",
+            headers,
+            body: body,
+            signal: AbortSignal.timeout(this.timeout),
+        });
+        if (!response.ok) {
+            throw await this.parseError(response);
+        }
+        const data = (await response.json());
+        return {
+            results: data.results,
+            query: data.query,
+            searchTime: data.search_time,
+            originalInputTokens: data.original_input_tokens,
+            outputTokens: data.output_tokens,
+            tokensSaved: data.original_input_tokens - data.output_tokens,
+        };
+    }
     async parseError(response) {
         let msg;
         try {

package/dist/compress.d.ts

@@ -30,6 +30,7 @@ export declare function compressOpenAIMessages(ttc: Compressor, messages: OpenAI
 interface AnthropicBlock {
     type: string;
     text?: string;
+    tool_use_id?: string;
     content?: string | Array<{
         type: string;
         text?: string;
@@ -42,7 +43,10 @@ interface AnthropicMessage {
     content?: string | AnthropicBlock[] | null;
     [key: string]: unknown;
 }
-export declare function compressAnthropicMessages(ttc: Compressor, messages: AnthropicMessage[], model: string, roleAggr: Record<string, number>): Promise<AnthropicMessage[]>;
+export declare function compressAnthropicMessages(ttc: Compressor, messages: AnthropicMessage[], model: string, roleAggr: Record<string, number>, options?: {
+    stripServerToolResults?: boolean;
+    skipToolName?: string;
+}): Promise<AnthropicMessage[]>;
 interface AISDKTextPart {
     type: "text";
     text: string;

package/dist/compress.js

@@ -1,4 +1,5 @@
 const DEFAULT_ROLES = ["user", "system", "tool"];
+const SERVER_TOOL_BLOCK_TYPES = new Set(["web_search_tool_result", "server_tool_use"]);
 export class StatsTTC {
     inner;
     stats;
@@ -46,8 +47,37 @@ export async function compressOpenAIMessages(ttc, messages, model, roleAggr) {
         return msg;
     }));
 }
-export async function compressAnthropicMessages(ttc, messages, model, roleAggr) {
+function collectToolUseIds(messages, toolName) {
+    const ids = new Set();
+    for (const msg of messages) {
+        if (msg.role !== "assistant" || !Array.isArray(msg.content))
+            continue;
+        for (const block of msg.content) {
+            if (block.type === "tool_use" && block.name === toolName && block.id) {
+                ids.add(block.id);
+            }
+        }
+    }
+    return ids;
+}
+export async function compressAnthropicMessages(ttc, messages, model, roleAggr, options) {
+    const stripServerToolResults = options?.stripServerToolResults ?? false;
+    const skipIds = options?.skipToolName ? collectToolUseIds(messages, options.skipToolName) : undefined;
     return Promise.all(messages.map(async (msg) => {
+        if (msg.role === "assistant") {
+            const assistantAggr = roleAggr["assistant"];
+            if (assistantAggr == null && !stripServerToolResults)
+                return msg;
+            if (typeof msg.content === "string" && assistantAggr != null && msg.content.trim()) {
+                const result = await ttc.compress(msg.content, { model, aggressiveness: assistantAggr });
+                return { ...msg, content: result.output };
+            }
+            if (Array.isArray(msg.content)) {
+                const blocks = await compressAssistantBlocks(ttc, msg.content, model, assistantAggr, stripServerToolResults);
+                return { ...msg, content: blocks };
+            }
+            return msg;
+        }
         if (msg.role !== "user")
             return msg;
         const userAggr = roleAggr["user"];
@@ -59,19 +89,35 @@ export async function compressAnthropicMessages(ttc, messages, model, roleAggr)
             return { ...msg, content: result.output };
         }
         if (Array.isArray(msg.content)) {
-            const blocks = await compressAnthropicBlocks(ttc, msg.content, model, userAggr, toolAggr);
+            const blocks = await compressAnthropicBlocks(ttc, msg.content, model, userAggr, toolAggr, skipIds);
             return { ...msg, content: blocks };
         }
         return msg;
     }));
 }
-async function compressAnthropicBlocks(ttc, blocks, model, userAggr, toolAggr) {
+async function compressAssistantBlocks(ttc, blocks, model, assistantAggr, stripServerToolResults) {
+    const results = await Promise.all(blocks.map(async (block) => {
+        if (stripServerToolResults && SERVER_TOOL_BLOCK_TYPES.has(block.type)) {
+            return null;
+        }
+        if (block.type === "text" && assistantAggr != null && typeof block.text === "string" && block.text.trim()) {
+            const result = await ttc.compress(block.text, { model, aggressiveness: assistantAggr });
+            return { ...block, text: result.output };
+        }
+        return block;
+    }));
+    const filtered = results.filter((b) => b != null);
+    return filtered.length > 0 ? filtered : blocks;
+}
+async function compressAnthropicBlocks(ttc, blocks, model, userAggr, toolAggr, skipIds) {
     return Promise.all(blocks.map(async (block) => {
         if (block.type === "text" && userAggr != null && typeof block.text === "string" && block.text.trim()) {
             const result = await ttc.compress(block.text, { model, aggressiveness: userAggr });
             return { ...block, text: result.output };
         }
         if (block.type === "tool_result" && toolAggr != null) {
+            if (block.tool_use_id && skipIds?.has(block.tool_use_id))
+                return block;
             return compressToolResult(ttc, block, model, toolAggr);
         }
         return block;

package/dist/index.d.ts

@@ -1,7 +1,7 @@
 export { TheTokenCompany, BEAR_1, BEAR_1_1, BEAR_1_2, BEAR_2 } from "./client.js";
 export { TheTokenCompanyError, AuthenticationError, InvalidRequestError, PaymentRequiredError, RequestTooLargeError, RateLimitError, APIError, } from "./errors.js";
 export { CompressionStats } from "./types.js";
-export type { CompressResult, TheTokenCompanyOptions, Aggressiveness, WithCompressionOptions, TurnStats, } from "./types.js";
+export type { CompressResult, TheTokenCompanyOptions, Aggressiveness, WithCompressionOptions, TurnStats, SearchRequestOptions, SearchResultItem, SearchResult, } from "./types.js";
 /**
  * Wrap text in `<ttc_safe>` tags to protect it from compression.
  */

package/dist/types.d.ts

@@ -31,9 +31,38 @@ export interface WithCompressionOptions {
     compressionApiKey: string;
     model?: string;
     aggressiveness?: Aggressiveness;
+    /** Compress text blocks in assistant messages (multi-turn web search optimization). */
+    compressAssistant?: boolean;
+    /** Strip server-side tool result blocks (e.g. web_search_tool_result) from assistant messages. Disables citations in subsequent turns. */
+    stripServerToolResults?: boolean;
+    /** Replace Anthropic's server-side web search with TTC's compressible search. */
+    webSearch?: boolean;
+    baseUrl?: string;
     appId?: string;
     fetch?: typeof globalThis.fetch;
 }
+export interface SearchRequestOptions {
+    maxResults?: number;
+    searchDepth?: "basic" | "advanced";
+    includeRawContent?: boolean;
+    model?: string;
+    aggressiveness?: number;
+    appId?: string;
+}
+export interface SearchResultItem {
+    url: string;
+    title: string;
+    content: string;
+    score?: number;
+}
+export interface SearchResult {
+    results: SearchResultItem[];
+    query: string;
+    searchTime: number;
+    originalInputTokens: number;
+    outputTokens: number;
+    tokensSaved: number;
+}
 export interface TurnStats {
     inputTokens: number;
     outputTokens: number;

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "the-token-company",
-  "version": "0.2.5",
+  "version": "0.3.1",
   "description": "Node.js SDK for The Token Company — compress LLM prompts to reduce costs and latency",
   "type": "module",
   "main": "./dist/index.js",