@illuma-ai/agents 1.0.96 → 1.1.0

package/src/graphs/gapFeatures.test.ts

@@ -0,0 +1,520 @@
+/**
+ * gapFeatures.test.ts
+ *
+ * Integration tests for the four LibreChat gap features:
+ * 1. Tool Discovery Caching
+ * 2. SummarizationConfig (trigger types, initialSummary)
+ * 3. EMA Pruning Calibration
+ * 4. Message Deduplication
+ *
+ * These tests verify the features work together in the Graph pipeline
+ * without breaking existing functionality.
+ */
+
+import {
+  HumanMessage,
+  AIMessage,
+  AIMessageChunk,
+  SystemMessage,
+  ToolMessage,
+  BaseMessage,
+} from '@langchain/core/messages';
+import type { TokenCounter } from '@/types/run';
+import type { SummarizationConfig } from '@/types/graph';
+import { createPruneMessages } from '@/messages/prune';
+import { deduplicateSystemMessages } from '@/messages/dedup';
+import { ToolDiscoveryCache } from '@/utils/toolDiscoveryCache';
+import {
+  createPruneCalibration,
+  updatePruneCalibration,
+  applyCalibration,
+} from '@/utils/pruneCalibration';
+import { Constants } from '@/common';
+
+const simpleTokenCounter: TokenCounter = (msg: BaseMessage): number => {
+  const content =
+    typeof msg.content === 'string' ? msg.content : JSON.stringify(msg.content);
+  return Math.ceil(content.length / 4);
+};
+
+function buildConversation(
+  messageCount: number,
+  charsPerMsg = 400
+): BaseMessage[] {
+  const messages: BaseMessage[] = [
+    new SystemMessage('You are a helpful assistant.'),
+  ];
+  for (let i = 0; i < messageCount; i++) {
+    const text = `Message ${i}: ${'x'.repeat(charsPerMsg - 15)}`;
+    messages.push(i % 2 === 0 ? new HumanMessage(text) : new AIMessage(text));
+  }
+  return messages;
+}
+
+// ============================================================================
+// 1. Tool Discovery Caching Integration
+// ============================================================================
+
+describe('Tool Discovery Caching — Integration', () => {
+  it('caches tool discoveries across multiple pruning iterations', () => {
+    const cache = new ToolDiscoveryCache();
+    const messages: BaseMessage[] = [
+      new SystemMessage('System prompt'),
+      new HumanMessage('Find a tool'),
+      new AIMessageChunk({
+        content: 'Searching',
+        tool_calls: [
+          { id: 'tc_1', name: Constants.TOOL_SEARCH, args: { query: 'web' } },
+        ],
+      }),
+      new ToolMessage({
+        content: 'Found tools',
+        tool_call_id: 'tc_1',
+        name: Constants.TOOL_SEARCH,
+        artifact: { tool_references: [{ tool_name: 'web_search' }] },
+      }),
+    ];
+
+    // Iteration 1: discovers web_search
+    const disc1 = cache.getNewDiscoveries(messages);
+    expect(disc1).toEqual(['web_search']);
+
+    // Iteration 2: same messages, no new discoveries (cached)
+    const disc2 = cache.getNewDiscoveries(messages);
+    expect(disc2).toEqual([]);
+
+    // Iteration 3: new tool search added
+    messages.push(
+      new AIMessageChunk({
+        content: 'More tools',
+        tool_calls: [
+          { id: 'tc_2', name: Constants.TOOL_SEARCH, args: { query: 'code' } },
+        ],
+      }),
+      new ToolMessage({
+        content: 'Found more',
+        tool_call_id: 'tc_2',
+        name: Constants.TOOL_SEARCH,
+        artifact: {
+          tool_references: [
+            { tool_name: 'code_exec' },
+            { tool_name: 'web_search' },
+          ],
+        },
+      })
+    );
+
+    // Only code_exec is new (web_search already cached)
+    const disc3 = cache.getNewDiscoveries(messages);
+    expect(disc3).toEqual(['code_exec']);
+    expect(cache.size).toBe(2);
+  });
+
+  it('seed + incremental discovery simulates cross-turn caching', () => {
+    const cache = new ToolDiscoveryCache();
+
+    // Seed from prior turn's discoveries
+    cache.seed(['tool_a', 'tool_b']);
+
+    const messages: BaseMessage[] = [
+      new HumanMessage('Use tool_a and find tool_c'),
+      new AIMessageChunk({
+        content: 'Searching',
+        tool_calls: [{ id: 'tc_1', name: Constants.TOOL_SEARCH, args: {} }],
+      }),
+      new ToolMessage({
+        content: 'Found',
+        tool_call_id: 'tc_1',
+        name: Constants.TOOL_SEARCH,
+        artifact: {
+          tool_references: [{ tool_name: 'tool_a' }, { tool_name: 'tool_c' }],
+        },
+      }),
+    ];
+
+    const newDisc = cache.getNewDiscoveries(messages);
+    // tool_a is seeded, only tool_c is new
+    expect(newDisc).toEqual(['tool_c']);
+    expect(cache.getAllDiscoveredTools()).toEqual(
+      expect.arrayContaining(['tool_a', 'tool_b', 'tool_c'])
+    );
+  });
+});
+
+// ============================================================================
+// 2. SummarizationConfig Integration
+// ============================================================================
+
+describe('SummarizationConfig — Trigger Logic', () => {
+  /**
+   * Simulates the Graph's shouldTriggerSummarization logic.
+   */
+  function shouldTriggerSummarization(
+    prunedMessageCount: number,
+    maxContextTokens: number,
+    indexTokenCountMap: Record<string, number | undefined>,
+    instructionTokens: number,
+    config?: SummarizationConfig
+  ): boolean {
+    if (prunedMessageCount === 0) return false;
+    if (!config || !config.triggerType) return true;
+
+    const threshold = config.triggerThreshold;
+
+    switch (config.triggerType) {
+      case 'contextPercentage': {
+        if (maxContextTokens <= 0) return true;
+        const effectiveThreshold = threshold ?? 80;
+        let totalTokens = instructionTokens;
+        for (const key in indexTokenCountMap) {
+          totalTokens += indexTokenCountMap[key] ?? 0;
+        }
+        const utilization = (totalTokens / maxContextTokens) * 100;
+        return utilization >= effectiveThreshold;
+      }
+      case 'messageCount': {
+        const effectiveThreshold = threshold ?? 5;
+        return prunedMessageCount >= effectiveThreshold;
+      }
+      case 'tokenThreshold': {
+        if (threshold == null) return true;
+        let totalTokens = instructionTokens;
+        for (const key in indexTokenCountMap) {
+          totalTokens += indexTokenCountMap[key] ?? 0;
+        }
+        return totalTokens >= threshold;
+      }
+      default:
+        return true;
+    }
+  }
+
+  it('no config = always triggers (backward compatible)', () => {
+    expect(shouldTriggerSummarization(3, 10000, {}, 100)).toBe(true);
+  });
+
+  it('contextPercentage: triggers at 80% utilization', () => {
+    const tokenMap = { '0': 4000, '1': 3500, '2': 500 };
+    // Total = 100 + 8000 = 8100, 8100/10000 = 81% > 80%
+    expect(
+      shouldTriggerSummarization(2, 10000, tokenMap, 100, {
+        triggerType: 'contextPercentage',
+        triggerThreshold: 80,
+      })
+    ).toBe(true);
+  });
+
+  it('contextPercentage: does NOT trigger below threshold', () => {
+    const tokenMap = { '0': 2000, '1': 1000 };
+    // Total = 100 + 3000 = 3100, 3100/10000 = 31% < 80%
+    expect(
+      shouldTriggerSummarization(2, 10000, tokenMap, 100, {
+        triggerType: 'contextPercentage',
+        triggerThreshold: 80,
+      })
+    ).toBe(false);
+  });
+
+  it('messageCount: triggers when enough messages pruned', () => {
+    expect(
+      shouldTriggerSummarization(5, 10000, {}, 100, {
+        triggerType: 'messageCount',
+        triggerThreshold: 5,
+      })
+    ).toBe(true);
+
+    expect(
+      shouldTriggerSummarization(3, 10000, {}, 100, {
+        triggerType: 'messageCount',
+        triggerThreshold: 5,
+      })
+    ).toBe(false);
+  });
+
+  it('tokenThreshold: triggers when total tokens exceed threshold', () => {
+    const tokenMap = { '0': 5000, '1': 3000 };
+    expect(
+      shouldTriggerSummarization(2, 10000, tokenMap, 100, {
+        triggerType: 'tokenThreshold',
+        triggerThreshold: 8000,
+      })
+    ).toBe(true);
+
+    expect(
+      shouldTriggerSummarization(2, 10000, tokenMap, 100, {
+        triggerType: 'tokenThreshold',
+        triggerThreshold: 9000,
+      })
+    ).toBe(false);
+  });
+
+  it('never triggers with 0 pruned messages', () => {
+    expect(shouldTriggerSummarization(0, 10000, {}, 100)).toBe(false);
+    expect(
+      shouldTriggerSummarization(0, 10000, {}, 100, {
+        triggerType: 'messageCount',
+        triggerThreshold: 1,
+      })
+    ).toBe(false);
+  });
+
+  it('initialSummary provides cross-run seeding', () => {
+    const config: SummarizationConfig = {
+      initialSummary: 'This agent helps with data analysis tasks.',
+    };
+
+    // Simulate the Graph logic: when no cached/persisted summary exists,
+    // initialSummary is used as fallback
+    let summary: string | undefined;
+    const cachedRunSummary: string | null = null;
+    const persistedSummary: string | null = null;
+
+    if (cachedRunSummary != null) {
+      summary = cachedRunSummary;
+    } else if (persistedSummary != null && persistedSummary !== '') {
+      summary = persistedSummary;
+    } else if (config.initialSummary != null && config.initialSummary !== '') {
+      summary = config.initialSummary;
+    }
+
+    expect(summary).toBe('This agent helps with data analysis tasks.');
+  });
+});
+
+// ============================================================================
+// 3. EMA Pruning Calibration Integration
+// ============================================================================
+
+describe('EMA Pruning Calibration — Integration', () => {
+  it('adjusts pruning budget across iterations', () => {
+    let calibration = createPruneCalibration();
+    const rawBudget = 10000;
+
+    // Iteration 1: no calibration data → raw budget
+    expect(applyCalibration(rawBudget, calibration)).toBe(10000);
+
+    // Simulate: our counter estimates 8000 tokens but API says 10000
+    // (we're under-counting → need to prune more aggressively)
+    calibration = updatePruneCalibration(calibration, 10000, 8000);
+    const adjusted = applyCalibration(rawBudget, calibration);
+    expect(adjusted).toBeLessThan(10000); // More aggressive pruning
+  });
+
+  it('full pruning cycle with calibration', () => {
+    let calibration = createPruneCalibration();
+    const messages = buildConversation(40);
+
+    // Iteration 1: uncalibrated
+    const maxTokens1 = applyCalibration(200, calibration);
+    const prune1 = createPruneMessages({
+      startIndex: 0,
+      provider: 'anthropic' as any,
+      tokenCounter: simpleTokenCounter,
+      maxTokens: maxTokens1,
+      indexTokenCountMap: {},
+    });
+    const result1 = prune1({ messages });
+    expect(result1.messagesToRefine.length).toBeGreaterThan(0);
+
+    // Simulate API returning higher token count than our estimate
+    calibration = updatePruneCalibration(calibration, 250, 200);
+
+    // Iteration 2: calibrated (should use adjusted budget)
+    const maxTokens2 = applyCalibration(200, calibration);
+    expect(maxTokens2).not.toBe(200); // Budget adjusted
+
+    // Multiple iterations should converge
+    for (let i = 0; i < 5; i++) {
+      calibration = updatePruneCalibration(calibration, 250, 200);
+    }
+    const finalBudget = applyCalibration(200, calibration);
+    // Should stabilize around 200 * (200/250) ≈ 160
+    expect(finalBudget).toBeLessThan(200);
+    expect(finalBudget).toBeGreaterThan(100);
+  });
+});
+
+// ============================================================================
+// 4. Message Deduplication Integration
+// ============================================================================
+
+describe('Message Deduplication — Integration', () => {
+  it('deduplicates post-prune notes from multiple iterations', () => {
+    const postPruneNote = 'Note: Earlier messages have been compressed.';
+    const messages: BaseMessage[] = [
+      new SystemMessage('Main system prompt'),
+      new SystemMessage('[Conversation Summary]\nPrior context'),
+      new SystemMessage(postPruneNote), // Iteration 1 post-prune note
+      new HumanMessage('Q1'),
+      new AIMessage('A1'),
+      new SystemMessage(postPruneNote), // Iteration 2 duplicate
+      new HumanMessage('Q2'),
+      new AIMessage('A2'),
+      new SystemMessage(postPruneNote), // Iteration 3 duplicate
+    ];
+
+    const { messages: deduped, removedCount } =
+      deduplicateSystemMessages(messages);
+    expect(removedCount).toBe(2);
+    expect(deduped).toHaveLength(7);
+
+    // Non-system messages all preserved
+    const humanMsgs = deduped.filter((m) => m.getType() === 'human');
+    expect(humanMsgs).toHaveLength(2);
+  });
+
+  it('preserves unique system messages including summary', () => {
+    const messages: BaseMessage[] = [
+      new SystemMessage('Main prompt'),
+      new SystemMessage('[Conversation Summary]\nVersion 1'),
+      new HumanMessage('Q'),
+      new SystemMessage('[Conversation Summary]\nVersion 2 - updated'),
+    ];
+
+    const { messages: deduped, removedCount } =
+      deduplicateSystemMessages(messages);
+    expect(removedCount).toBe(0);
+    expect(deduped).toHaveLength(4);
+  });
+
+  it('works with pruning + dedup pipeline', () => {
+    // Simulate: prune messages, inject summary, then dedup
+    const allMessages = buildConversation(20);
+
+    // Step 1: Prune
+    const prune = createPruneMessages({
+      startIndex: 0,
+      provider: 'anthropic' as any,
+      tokenCounter: simpleTokenCounter,
+      maxTokens: 300,
+      indexTokenCountMap: {},
+    });
+    const { context, messagesToRefine } = prune({ messages: allMessages });
+    expect(messagesToRefine.length).toBeGreaterThan(0);
+
+    // Step 2: Inject summary
+    const summaryMsg = new SystemMessage(
+      '[Conversation Summary]\nUser discussed tasks'
+    );
+    const systemIdx = context[0]?.getType() === 'system' ? 1 : 0;
+    let withSummary = [
+      ...context.slice(0, systemIdx),
+      summaryMsg,
+      ...context.slice(systemIdx),
+    ];
+
+    // Simulate adding post-prune note
+    withSummary = [...withSummary, new SystemMessage('Context was compressed')];
+
+    // Step 3: Dedup (should not remove anything since all unique)
+    const { messages: final, removedCount } =
+      deduplicateSystemMessages(withSummary);
+    expect(removedCount).toBe(0);
+    expect(final.length).toBe(withSummary.length);
+  });
+});
+
+// ============================================================================
+// Combined Integration
+// ============================================================================
+
+describe('All Features Combined — Full Pipeline', () => {
+  it('simulates 3-turn conversation with all features active', async () => {
+    const toolCache = new ToolDiscoveryCache();
+    let calibration = createPruneCalibration();
+    let persistedSummary: string | null = null;
+    const sumConfig: SummarizationConfig = {
+      triggerType: 'contextPercentage',
+      triggerThreshold: 50,
+      reserveRatio: 0.3,
+    };
+
+    const callback = jest
+      .fn()
+      .mockImplementation(async (msgs: BaseMessage[]) => {
+        const summary = `Summary of ${msgs.length} messages`;
+        persistedSummary = summary;
+        return summary;
+      });
+
+    for (let turn = 0; turn < 3; turn++) {
+      // Build conversation that exceeds budget
+      const messages = buildConversation(15);
+
+      // Tool discovery (turn 1 has tool search results)
+      if (turn === 0) {
+        messages.push(
+          new AIMessageChunk({
+            content: 'Searching',
+            tool_calls: [
+              { id: `tc_${turn}`, name: Constants.TOOL_SEARCH, args: {} },
+            ],
+          }),
+          new ToolMessage({
+            content: 'Found',
+            tool_call_id: `tc_${turn}`,
+            name: Constants.TOOL_SEARCH,
+            artifact: { tool_references: [{ tool_name: 'web_search' }] },
+          })
+        );
+      }
+
+      const discoveries = toolCache.getNewDiscoveries(messages);
+      if (turn === 0) {
+        expect(discoveries).toEqual(['web_search']);
+      } else {
+        expect(discoveries).toEqual([]);
+      }
+
+      // Prune with calibration
+      const maxTokens = applyCalibration(300, calibration);
+      const prune = createPruneMessages({
+        startIndex: 0,
+        provider: 'anthropic' as any,
+        tokenCounter: simpleTokenCounter,
+        maxTokens,
+        indexTokenCountMap: {},
+      });
+      const { context, messagesToRefine } = prune({ messages });
+
+      let assembled = [...context];
+
+      // Inject summary if available
+      if (persistedSummary && messagesToRefine.length > 0) {
+        const summaryMsg = new SystemMessage(
+          `[Conversation Summary]\n${persistedSummary}`
+        );
+        const sysIdx = assembled[0]?.getType() === 'system' ? 1 : 0;
+        assembled = [
+          ...assembled.slice(0, sysIdx),
+          summaryMsg,
+          ...assembled.slice(sysIdx),
+        ];
+      }
+
+      // Fire background summary
+      if (messagesToRefine.length > 0) {
+        callback(messagesToRefine).catch(() => {});
+        await new Promise((r) => setTimeout(r, 10));
+      }
+
+      // Dedup
+      const { messages: deduped } = deduplicateSystemMessages(assembled);
+      expect(deduped.length).toBeLessThanOrEqual(assembled.length);
+
+      // Update calibration (simulated API response)
+      calibration = updatePruneCalibration(
+        calibration,
+        maxTokens + 50,
+        maxTokens
+      );
+    }
+
+    // Verify state after 3 turns
+    expect(toolCache.size).toBe(1);
+    expect(toolCache.has('web_search')).toBe(true);
+    expect(calibration.iterations).toBe(3);
+    expect(persistedSummary).toContain('Summary of');
+    expect(callback).toHaveBeenCalled();
+  });
+});