@illuma-ai/agents 1.0.96 → 1.1.0

package/src/utils/pruneCalibration.ts

@@ -0,0 +1,92 @@
+// src/utils/pruneCalibration.ts
+import type { PruneCalibrationState } from '@/types/graph';
+import {
+  PRUNING_EMA_ALPHA,
+  PRUNING_INITIAL_CALIBRATION,
+} from '@/common/constants';
+
+/**
+ * Creates an initial pruning calibration state.
+ *
+ * @param initialRatio - Starting calibration ratio (default: 1.0)
+ * @returns Fresh calibration state
+ */
+export function createPruneCalibration(
+  initialRatio?: number
+): PruneCalibrationState {
+  return {
+    ratio: initialRatio ?? PRUNING_INITIAL_CALIBRATION,
+    iterations: 0,
+  };
+}
+
+/**
+ * Updates the pruning calibration using Exponential Moving Average (EMA).
+ *
+ * Problem: Without calibration, the pruner's token estimates can diverge from
+ * reality across iterations, causing either:
+ * - Over-pruning (context cliff): Too many messages removed at once, losing critical tool results
+ * - Under-pruning: Not enough messages removed, hitting hard token limits
+ *
+ * Solution: Track the ratio between actual token usage (from API response) and
+ * estimated token usage (from our token counter). Apply EMA smoothing so the
+ * calibration adjusts gradually, preventing oscillation.
+ *
+ * The calibration ratio is applied to maxTokens in the pruner:
+ *   effectiveMaxTokens = maxTokens * calibrationRatio
+ *
+ * If actual > estimated → ratio decreases → prune more aggressively
+ * If actual < estimated → ratio increases → prune less aggressively
+ *
+ * @param state - Current calibration state
+ * @param actualTokens - Actual token count from API response (UsageMetadata)
+ * @param estimatedTokens - Estimated token count from token counter
+ * @param alpha - EMA smoothing factor (default: PRUNING_EMA_ALPHA)
+ * @returns Updated calibration state (new object, does not mutate input)
+ */
+export function updatePruneCalibration(
+  state: PruneCalibrationState,
+  actualTokens: number,
+  estimatedTokens: number,
+  alpha: number = PRUNING_EMA_ALPHA
+): PruneCalibrationState {
+  // Guard against division by zero or invalid inputs
+  if (estimatedTokens <= 0 || actualTokens <= 0) {
+    return state;
+  }
+
+  // Raw ratio: how much our estimate differs from reality
+  const observedRatio = estimatedTokens / actualTokens;
+
+  // Clamp to prevent extreme adjustments from outlier readings
+  // Range [0.5, 2.0] means we never more than double or halve the budget
+  const clampedRatio = Math.max(0.5, Math.min(2.0, observedRatio));
+
+  // Apply EMA: new_ratio = α * observed + (1 - α) * previous
+  const newRatio = alpha * clampedRatio + (1 - alpha) * state.ratio;
+
+  return {
+    ratio: newRatio,
+    iterations: state.iterations + 1,
+  };
+}
+
+/**
+ * Applies the calibration ratio to a max token budget.
+ * The ratio adjusts the effective budget so pruning is more or less aggressive
+ * based on observed vs. estimated token divergence.
+ *
+ * @param maxTokens - Raw max token budget
+ * @param state - Current calibration state
+ * @returns Adjusted max token budget
+ */
+export function applyCalibration(
+  maxTokens: number,
+  state: PruneCalibrationState
+): number {
+  if (state.iterations === 0) {
+    // No calibration data yet — use raw budget
+    return maxTokens;
+  }
+  return Math.floor(maxTokens * state.ratio);
+}

package/src/utils/run.ts

@@ -1,108 +1,108 @@
-import { CallbackManagerForChainRun } from '@langchain/core/callbacks/manager';
-import {
-  mergeConfigs,
-  patchConfig,
-  Runnable,
-  RunnableConfig,
-} from '@langchain/core/runnables';
-import { AsyncLocalStorageProviderSingleton } from '@langchain/core/singletons';
-
-/**
- * Delays the execution for a specified number of milliseconds.
- *
- * @param {number} ms - The number of milliseconds to delay.
- * @return {Promise<void>} A promise that resolves after the specified delay.
- */
-export function sleep(ms: number): Promise<void> {
-  return new Promise((resolve) => setTimeout(resolve, ms));
-}
-
-// eslint-disable-next-line @typescript-eslint/no-explicit-any
-export interface RunnableCallableArgs extends Partial<any> {
-  name?: string;
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  func: (...args: any[]) => any;
-  tags?: string[];
-  trace?: boolean;
-  recurse?: boolean;
-}
-
-export class RunnableCallable<I = unknown, O = unknown> extends Runnable<I, O> {
-  lc_namespace: string[] = ['langgraph'];
-
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  func: (...args: any[]) => any;
-
-  tags?: string[];
-
-  config?: RunnableConfig;
-
-  trace: boolean = true;
-
-  recurse: boolean = true;
-
-  constructor(fields: RunnableCallableArgs) {
-    super();
-    this.name = fields.name ?? fields.func.name;
-    this.func = fields.func;
-    this.config = fields.tags ? { tags: fields.tags } : undefined;
-    this.trace = fields.trace ?? this.trace;
-    this.recurse = fields.recurse ?? this.recurse;
-  }
-
-  protected async _tracedInvoke(
-    input: I,
-    config?: Partial<RunnableConfig>,
-    runManager?: CallbackManagerForChainRun
-  ): Promise<O> {
-    return new Promise<O>((resolve, reject) => {
-      // Defensive check: ensure runManager has getChild method before calling
-      const childCallbacks =
-        typeof runManager?.getChild === 'function'
-          ? runManager.getChild()
-          : undefined;
-      let childConfig: Partial<RunnableConfig> | null = patchConfig(config, {
-        callbacks: childCallbacks,
-      });
-      void AsyncLocalStorageProviderSingleton.runWithConfig(
-        childConfig,
-        async () => {
-          try {
-            const output = await this.func(input, childConfig);
-            childConfig = null;
-            resolve(output);
-          } catch (e) {
-            childConfig = null;
-            reject(e);
-          }
-        }
-      );
-    });
-  }
-
-  async invoke(
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    input: any,
-    options?: Partial<RunnableConfig> | undefined
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  ): Promise<any> {
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    let returnValue: any;
-
-    if (this.trace) {
-      returnValue = await this._callWithConfig(
-        this._tracedInvoke,
-        input,
-        mergeConfigs(this.config, options)
-      );
-    } else {
-      returnValue = await this.func(input, mergeConfigs(this.config, options));
-    }
-
-    if (Runnable.isRunnable(returnValue) && this.recurse) {
-      return await returnValue.invoke(input, options);
-    }
-
-    return returnValue;
-  }
-}
+import { CallbackManagerForChainRun } from '@langchain/core/callbacks/manager';
+import {
+  mergeConfigs,
+  patchConfig,
+  Runnable,
+  RunnableConfig,
+} from '@langchain/core/runnables';
+import { AsyncLocalStorageProviderSingleton } from '@langchain/core/singletons';
+
+/**
+ * Delays the execution for a specified number of milliseconds.
+ *
+ * @param {number} ms - The number of milliseconds to delay.
+ * @return {Promise<void>} A promise that resolves after the specified delay.
+ */
+export function sleep(ms: number): Promise<void> {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}
+
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export interface RunnableCallableArgs extends Partial<any> {
+  name?: string;
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  func: (...args: any[]) => any;
+  tags?: string[];
+  trace?: boolean;
+  recurse?: boolean;
+}
+
+export class RunnableCallable<I = unknown, O = unknown> extends Runnable<I, O> {
+  lc_namespace: string[] = ['langgraph'];
+
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  func: (...args: any[]) => any;
+
+  tags?: string[];
+
+  config?: RunnableConfig;
+
+  trace: boolean = true;
+
+  recurse: boolean = true;
+
+  constructor(fields: RunnableCallableArgs) {
+    super();
+    this.name = fields.name ?? fields.func.name;
+    this.func = fields.func;
+    this.config = fields.tags ? { tags: fields.tags } : undefined;
+    this.trace = fields.trace ?? this.trace;
+    this.recurse = fields.recurse ?? this.recurse;
+  }
+
+  protected async _tracedInvoke(
+    input: I,
+    config?: Partial<RunnableConfig>,
+    runManager?: CallbackManagerForChainRun
+  ): Promise<O> {
+    return new Promise<O>((resolve, reject) => {
+      // Defensive check: ensure runManager has getChild method before calling
+      const childCallbacks =
+        typeof runManager?.getChild === 'function'
+          ? runManager.getChild()
+          : undefined;
+      let childConfig: Partial<RunnableConfig> | null = patchConfig(config, {
+        callbacks: childCallbacks,
+      });
+      void AsyncLocalStorageProviderSingleton.runWithConfig(
+        childConfig,
+        async () => {
+          try {
+            const output = await this.func(input, childConfig);
+            childConfig = null;
+            resolve(output);
+          } catch (e) {
+            childConfig = null;
+            reject(e);
+          }
+        }
+      );
+    });
+  }
+
+  async invoke(
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    input: any,
+    options?: Partial<RunnableConfig> | undefined
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  ): Promise<any> {
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    let returnValue: any;
+
+    if (this.trace) {
+      returnValue = await this._callWithConfig(
+        this._tracedInvoke,
+        input,
+        mergeConfigs(this.config, options)
+      );
+    } else {
+      returnValue = await this.func(input, mergeConfigs(this.config, options));
+    }
+
+    if (Runnable.isRunnable(returnValue) && this.recurse) {
+      return await returnValue.invoke(input, options);
+    }
+
+    return returnValue;
+  }
+}

package/src/utils/tokens.ts

@@ -1,118 +1,118 @@
-import { Tokenizer } from 'ai-tokenizer';
-import type { BaseMessage } from '@langchain/core/messages';
-import { ContentTypes } from '@/common/enum';
-
-export type EncodingName = 'o200k_base' | 'claude';
-
-const tokenizers: Partial<Record<EncodingName, Tokenizer>> = {};
-
-async function getTokenizer(
-  encoding: EncodingName = 'o200k_base'
-): Promise<Tokenizer> {
-  const cached = tokenizers[encoding];
-  if (cached) {
-    return cached;
-  }
-  const data =
-    encoding === 'claude'
-      ? await import('ai-tokenizer/encoding/claude')
-      : await import('ai-tokenizer/encoding/o200k_base');
-  const instance = new Tokenizer(data);
-  tokenizers[encoding] = instance;
-  return instance;
-}
-
-export function encodingForModel(model: string): EncodingName {
-  if (model.toLowerCase().includes('claude')) {
-    return 'claude';
-  }
-  return 'o200k_base';
-}
-
-export function getTokenCountForMessage(
-  message: BaseMessage,
-  getTokenCount: (text: string) => number
-): number {
-  const tokensPerMessage = 3;
-
-  const processValue = (value: unknown): void => {
-    if (Array.isArray(value)) {
-      for (const item of value) {
-        if (
-          !item ||
-          !item.type ||
-          item.type === ContentTypes.ERROR ||
-          item.type === ContentTypes.IMAGE_URL
-        ) {
-          continue;
-        }
-
-        if (item.type === ContentTypes.TOOL_CALL && item.tool_call != null) {
-          const toolName = item.tool_call?.name || '';
-          if (toolName != null && toolName && typeof toolName === 'string') {
-            numTokens += getTokenCount(toolName);
-          }
-
-          const args = item.tool_call?.args || '';
-          if (args != null && args && typeof args === 'string') {
-            numTokens += getTokenCount(args);
-          }
-
-          const output = item.tool_call?.output || '';
-          if (output != null && output && typeof output === 'string') {
-            numTokens += getTokenCount(output);
-          }
-          continue;
-        }
-
-        const nestedValue = item[item.type];
-
-        if (!nestedValue) {
-          continue;
-        }
-
-        processValue(nestedValue);
-      }
-    } else if (typeof value === 'string') {
-      numTokens += getTokenCount(value);
-    } else if (typeof value === 'number') {
-      numTokens += getTokenCount(value.toString());
-    } else if (typeof value === 'boolean') {
-      numTokens += getTokenCount(value.toString());
-    }
-  };
-
-  let numTokens = tokensPerMessage;
-  processValue(message.content);
-  return numTokens;
-}
-
-/**
- * Creates a token counter function using the specified encoding.
- * Lazily loads the encoding data on first use via dynamic import.
- */
-export const createTokenCounter = async (
-  encoding: EncodingName = 'o200k_base'
-): Promise<(message: BaseMessage) => number> => {
-  const tok = await getTokenizer(encoding);
-  const countTokens = (text: string): number => tok.count(text);
-  return (message: BaseMessage): number =>
-    getTokenCountForMessage(message, countTokens);
-};
-
-/** Utility to manage the token encoder lifecycle explicitly. */
-export const TokenEncoderManager = {
-  async initialize(): Promise<void> {
-    // No-op: ai-tokenizer is synchronously initialized from bundled data.
-  },
-
-  reset(): void {
-    for (const key of Object.keys(tokenizers)) {
-      delete tokenizers[key as EncodingName];
-    }
-  },
-
-  isInitialized(): boolean {
-    return Object.keys(tokenizers).length > 0;
-  },
-};
+import { Tokenizer } from 'ai-tokenizer';
+import type { BaseMessage } from '@langchain/core/messages';
+import { ContentTypes } from '@/common/enum';
+
+export type EncodingName = 'o200k_base' | 'claude';
+
+const tokenizers: Partial<Record<EncodingName, Tokenizer>> = {};
+
+async function getTokenizer(
+  encoding: EncodingName = 'o200k_base'
+): Promise<Tokenizer> {
+  const cached = tokenizers[encoding];
+  if (cached) {
+    return cached;
+  }
+  const data =
+    encoding === 'claude'
+      ? await import('ai-tokenizer/encoding/claude')
+      : await import('ai-tokenizer/encoding/o200k_base');
+  const instance = new Tokenizer(data);
+  tokenizers[encoding] = instance;
+  return instance;
+}
+
+export function encodingForModel(model: string): EncodingName {
+  if (model.toLowerCase().includes('claude')) {
+    return 'claude';
+  }
+  return 'o200k_base';
+}
+
+export function getTokenCountForMessage(
+  message: BaseMessage,
+  getTokenCount: (text: string) => number
+): number {
+  const tokensPerMessage = 3;
+
+  const processValue = (value: unknown): void => {
+    if (Array.isArray(value)) {
+      for (const item of value) {
+        if (
+          !item ||
+          !item.type ||
+          item.type === ContentTypes.ERROR ||
+          item.type === ContentTypes.IMAGE_URL
+        ) {
+          continue;
+        }
+
+        if (item.type === ContentTypes.TOOL_CALL && item.tool_call != null) {
+          const toolName = item.tool_call?.name || '';
+          if (toolName != null && toolName && typeof toolName === 'string') {
+            numTokens += getTokenCount(toolName);
+          }
+
+          const args = item.tool_call?.args || '';
+          if (args != null && args && typeof args === 'string') {
+            numTokens += getTokenCount(args);
+          }
+
+          const output = item.tool_call?.output || '';
+          if (output != null && output && typeof output === 'string') {
+            numTokens += getTokenCount(output);
+          }
+          continue;
+        }
+
+        const nestedValue = item[item.type];
+
+        if (!nestedValue) {
+          continue;
+        }
+
+        processValue(nestedValue);
+      }
+    } else if (typeof value === 'string') {
+      numTokens += getTokenCount(value);
+    } else if (typeof value === 'number') {
+      numTokens += getTokenCount(value.toString());
+    } else if (typeof value === 'boolean') {
+      numTokens += getTokenCount(value.toString());
+    }
+  };
+
+  let numTokens = tokensPerMessage;
+  processValue(message.content);
+  return numTokens;
+}
+
+/**
+ * Creates a token counter function using the specified encoding.
+ * Lazily loads the encoding data on first use via dynamic import.
+ */
+export const createTokenCounter = async (
+  encoding: EncodingName = 'o200k_base'
+): Promise<(message: BaseMessage) => number> => {
+  const tok = await getTokenizer(encoding);
+  const countTokens = (text: string): number => tok.count(text);
+  return (message: BaseMessage): number =>
+    getTokenCountForMessage(message, countTokens);
+};
+
+/** Utility to manage the token encoder lifecycle explicitly. */
+export const TokenEncoderManager = {
+  async initialize(): Promise<void> {
+    // No-op: ai-tokenizer is synchronously initialized from bundled data.
+  },
+
+  reset(): void {
+    for (const key of Object.keys(tokenizers)) {
+      delete tokenizers[key as EncodingName];
+    }
+  },
+
+  isInitialized(): boolean {
+    return Object.keys(tokenizers).length > 0;
+  },
+};