illuma-agents 1.0.2

package/src/specs/token-distribution-edge-case.test.ts

@@ -0,0 +1,316 @@
+// src/specs/token-distribution-edge-case.test.ts
+import {
+  HumanMessage,
+  AIMessage,
+  SystemMessage,
+  BaseMessage,
+} from '@langchain/core/messages';
+import type { UsageMetadata } from '@langchain/core/messages';
+import type * as t from '@/types';
+import { createPruneMessages } from '@/messages/prune';
+
+// Create a simple token counter for testing
+const createTestTokenCounter = (): t.TokenCounter => {
+  // This simple token counter just counts characters as tokens for predictable testing
+  return (message: BaseMessage): number => {
+    // Use type assertion to help TypeScript understand the type
+    const content = message.content as
+      | string
+      | Array<t.MessageContentComplex | string>
+      | undefined;
+
+    // Handle string content
+    if (typeof content === 'string') {
+      return content.length;
+    }
+
+    // Handle array content
+    if (Array.isArray(content)) {
+      let totalLength = 0;
+
+      for (const item of content) {
+        if (typeof item === 'string') {
+          totalLength += item.length;
+        } else if (typeof item === 'object') {
+          if ('text' in item && typeof item.text === 'string') {
+            totalLength += item.text.length;
+          }
+        }
+      }
+
+      return totalLength;
+    }
+
+    // Default case - if content is null, undefined, or any other type
+    return 0;
+  };
+};
+
+describe('Token Distribution Edge Case Tests', () => {
+  it('should only distribute tokens to messages that remain in the context after pruning', () => {
+    // Create a token counter
+    const tokenCounter = createTestTokenCounter();
+
+    // Create messages
+    const messages = [
+      new SystemMessage('System instruction'), // Will always be included
+      new HumanMessage('Message 1'), // Will be pruned
+      new AIMessage('Response 1'), // Will be pruned
+      new HumanMessage('Message 2'), // Will remain
+      new AIMessage('Response 2'), // Will remain
+    ];
+
+    // Calculate initial token counts for each message
+    const indexTokenCountMap: Record<string, number> = {
+      0: 17, // "System instruction"
+      1: 9, // "Message 1"
+      2: 10, // "Response 1"
+      3: 9, // "Message 2"
+      4: 10, // "Response 2"
+    };
+
+    // Set a token limit that will force pruning of the first two messages after the system message
+    const pruneMessages = createPruneMessages({
+      maxTokens: 40, // Only enough for system message + last two messages
+      startIndex: 0,
+      tokenCounter,
+      indexTokenCountMap: { ...indexTokenCountMap },
+    });
+
+    // First call to establish lastCutOffIndex
+    const initialResult = pruneMessages({ messages });
+
+    // Verify initial pruning
+    expect(initialResult.context.length).toBe(3);
+    expect(initialResult.context[0].content).toBe('System instruction');
+    expect(initialResult.context[1].content).toBe('Message 2');
+    expect(initialResult.context[2].content).toBe('Response 2');
+
+    // Now provide usage metadata with a different total token count
+    const usageMetadata: Partial<UsageMetadata> = {
+      input_tokens: 30,
+      output_tokens: 20,
+      total_tokens: 50, // Different from the sum of our initial token counts
+    };
+
+    // Call pruneMessages again with the usage metadata
+    const result = pruneMessages({
+      messages,
+      usageMetadata,
+    });
+
+    // The token distribution should only affect messages that remain in the context
+    // Messages at indices 0, 3, and 4 should have their token counts adjusted
+    // Messages at indices 1 and 2 should remain unchanged since they're pruned
+
+    // The token distribution should only affect messages that remain in the context
+    // Messages at indices 0, 3, and 4 should have their token counts adjusted
+    // Messages at indices 1 and 2 should remain unchanged since they're pruned
+
+    // Check that at least one of the pruned messages' token counts was not adjusted
+    // We're testing the principle that pruned messages don't get token redistribution
+    const atLeastOnePrunedMessageUnchanged =
+      result.indexTokenCountMap[1] === indexTokenCountMap[1] ||
+      result.indexTokenCountMap[2] === indexTokenCountMap[2];
+
+    expect(atLeastOnePrunedMessageUnchanged).toBe(true);
+
+    // Verify that the sum of tokens for messages in the context is close to the total_tokens from usageMetadata
+    // There might be small rounding differences or implementation details that affect the exact sum
+    const totalContextTokens =
+      (result.indexTokenCountMap[0] ?? 0) +
+      (result.indexTokenCountMap[3] ?? 0) +
+      (result.indexTokenCountMap[4] ?? 0);
+    expect(totalContextTokens).toBeGreaterThan(0);
+
+    // The key thing we're testing is that the token distribution happens for messages in the context
+    // and that the sum is reasonably close to the expected total
+    const tokenDifference = Math.abs(totalContextTokens - 50);
+    expect(tokenDifference).toBeLessThan(20); // Allow for some difference due to implementation details
+  });
+
+  it('should handle the case when all messages fit within the token limit', () => {
+    // Create a token counter
+    const tokenCounter = createTestTokenCounter();
+
+    // Create messages
+    const messages = [
+      new SystemMessage('System instruction'),
+      new HumanMessage('Message 1'),
+      new AIMessage('Response 1'),
+    ];
+
+    // Calculate initial token counts for each message
+    const indexTokenCountMap: Record<string, number> = {
+      0: 17, // "System instruction"
+      1: 9, // "Message 1"
+      2: 10, // "Response 1"
+    };
+
+    // Set a token limit that will allow all messages to fit
+    const pruneMessages = createPruneMessages({
+      maxTokens: 100,
+      startIndex: 0,
+      tokenCounter,
+      indexTokenCountMap: { ...indexTokenCountMap },
+    });
+
+    // First call to establish lastCutOffIndex (should be 0 since no pruning occurs)
+    const initialResult = pruneMessages({ messages });
+
+    // Verify no pruning occurred
+    expect(initialResult.context.length).toBe(3);
+
+    // Now provide usage metadata with a different total token count
+    const usageMetadata: Partial<UsageMetadata> = {
+      input_tokens: 20,
+      output_tokens: 10,
+      total_tokens: 30, // Different from the sum of our initial token counts
+    };
+
+    // Call pruneMessages again with the usage metadata
+    const result = pruneMessages({
+      messages,
+      usageMetadata,
+    });
+
+    // Since all messages fit, all token counts should be adjusted
+    const initialTotalTokens =
+      indexTokenCountMap[0] + indexTokenCountMap[1] + indexTokenCountMap[2];
+    const expectedRatio = 30 / initialTotalTokens;
+
+    // Check that all token counts were adjusted
+    expect(result.indexTokenCountMap[0]).toBe(
+      Math.round(indexTokenCountMap[0] * expectedRatio)
+    );
+    expect(result.indexTokenCountMap[1]).toBe(
+      Math.round(indexTokenCountMap[1] * expectedRatio)
+    );
+    expect(result.indexTokenCountMap[2]).toBe(
+      Math.round(indexTokenCountMap[2] * expectedRatio)
+    );
+
+    // Verify that the sum of all tokens equals the total_tokens from usageMetadata
+    const totalTokens =
+      (result.indexTokenCountMap[0] ?? 0) +
+      (result.indexTokenCountMap[1] ?? 0) +
+      (result.indexTokenCountMap[2] ?? 0);
+    expect(totalTokens).toBe(30);
+  });
+
+  it('should handle multiple pruning operations with token redistribution', () => {
+    // Create a token counter
+    const tokenCounter = createTestTokenCounter();
+
+    // Create a longer sequence of messages
+    const messages = [
+      new SystemMessage('System instruction'), // Will always be included
+      new HumanMessage('Message 1'), // Will be pruned in first round
+      new AIMessage('Response 1'), // Will be pruned in first round
+      new HumanMessage('Message 2'), // Will be pruned in second round
+      new AIMessage('Response 2'), // Will be pruned in second round
+      new HumanMessage('Message 3'), // Will remain
+      new AIMessage('Response 3'), // Will remain
+    ];
+
+    // Calculate initial token counts for each message
+    const indexTokenCountMap: Record<string, number> = {
+      0: 17, // "System instruction"
+      1: 9, // "Message 1"
+      2: 10, // "Response 1"
+      3: 9, // "Message 2"
+      4: 10, // "Response 2"
+      5: 9, // "Message 3"
+      6: 10, // "Response 3"
+    };
+
+    // Set a token limit that will force pruning
+    const pruneMessages = createPruneMessages({
+      maxTokens: 40, // Only enough for system message + last two messages
+      startIndex: 0,
+      tokenCounter,
+      indexTokenCountMap: { ...indexTokenCountMap },
+    });
+
+    // First pruning operation
+    const firstResult = pruneMessages({ messages });
+
+    // Verify first pruning
+    expect(firstResult.context.length).toBe(3);
+    expect(firstResult.context[0].content).toBe('System instruction');
+    expect(firstResult.context[1].content).toBe('Message 3');
+    expect(firstResult.context[2].content).toBe('Response 3');
+
+    // First usage metadata update
+    const firstUsageMetadata: Partial<UsageMetadata> = {
+      input_tokens: 30,
+      output_tokens: 20,
+      total_tokens: 50,
+    };
+
+    // Apply first usage metadata
+    const secondResult = pruneMessages({
+      messages,
+      usageMetadata: firstUsageMetadata,
+    });
+
+    // Add two more messages
+    messages.push(new HumanMessage('Message 4'));
+    const extendedMessages = [...messages, new AIMessage('Response 4')];
+
+    // Second usage metadata update
+    const secondUsageMetadata: Partial<UsageMetadata> = {
+      input_tokens: 30,
+      output_tokens: 20,
+      total_tokens: 50,
+    };
+
+    // Apply second usage metadata with extended messages
+    const thirdResult = pruneMessages({
+      messages: extendedMessages,
+      usageMetadata: secondUsageMetadata,
+    });
+
+    // The context should include the system message and some of the latest messages
+    expect(thirdResult.context.length).toBeGreaterThan(0);
+    expect(thirdResult.context[0].content).toBe('System instruction');
+    expect(thirdResult.context[1].content).toBe('Response 4');
+
+    // Find which messages are in the final context
+    const contextMessageIndices = thirdResult.context.map((msg) => {
+      // Find the index of this message in the original array
+      return extendedMessages.findIndex((m) => m.content === msg.content);
+    });
+
+    // Get the sum of token counts for messages in the context
+    let totalContextTokens = 0;
+    for (const idx of contextMessageIndices) {
+      totalContextTokens += thirdResult.indexTokenCountMap[idx] ?? 0;
+    }
+
+    // Verify that the sum of tokens for messages in the context is close to the total_tokens from usageMetadata
+    // There might be small rounding differences or implementation details that affect the exact sum
+    expect(totalContextTokens).toBeGreaterThan(0);
+
+    // The key thing we're testing is that the token distribution happens for messages in the context
+    // and that the sum is reasonably close to the expected total
+    const tokenDifference = Math.abs(totalContextTokens - 70);
+    expect(tokenDifference).toBeLessThan(50); // Allow for some difference due to implementation details
+
+    // Verify that messages not in the context have their original token counts or previously adjusted values
+    for (let i = 0; i < extendedMessages.length; i++) {
+      if (!contextMessageIndices.includes(i)) {
+        const expectedValue =
+          i < messages.length
+            ? (secondResult.indexTokenCountMap[i] ?? 0) || indexTokenCountMap[i]
+            : ((indexTokenCountMap as Record<string, number | undefined>)[i] ??
+              0);
+
+        const difference = Math.abs(
+          (thirdResult.indexTokenCountMap[i] ?? 0) - expectedValue
+        );
+        expect(difference).toBe(0);
+      }
+    }
+  });
+});

package/src/specs/tool-error.test.ts

@@ -0,0 +1,193 @@
+import { z } from 'zod';
+import { config } from 'dotenv';
+config();
+import { tool } from '@langchain/core/tools';
+import { ToolCall } from '@langchain/core/messages/tool';
+import { HumanMessage, BaseMessage } from '@langchain/core/messages';
+import type { RunnableConfig } from '@langchain/core/runnables';
+import type * as t from '@/types';
+import { ChatModelStreamHandler, createContentAggregator } from '@/stream';
+import { ToolEndHandler, ModelEndHandler } from '@/events';
+import { GraphEvents, Providers } from '@/common';
+import { getLLMConfig } from '@/utils/llmConfig';
+import { getArgs } from '@/scripts/args';
+import { StandardGraph } from '@/graphs';
+import { Run } from '@/run';
+
+const errorTool = tool(
+  async () => {
+    throw new Error('this is a test error I threw on purpose');
+  },
+  {
+    name: 'errorTool',
+    description: 'A tool that always throws an error',
+    schema: z.object({ input: z.string().optional() }),
+  }
+);
+
+describe('Tool Error Handling Tests', () => {
+  jest.setTimeout(30000);
+  let run: Run<t.IState>;
+  let contentParts: t.MessageContentComplex[];
+  let conversationHistory: BaseMessage[];
+  let aggregateContent: t.ContentAggregator;
+  let handleToolCallErrorSpy: jest.SpyInstance;
+
+  const config: Partial<RunnableConfig> & {
+    version: 'v1' | 'v2';
+    run_id?: string;
+    streamMode: string;
+  } = {
+    configurable: {
+      thread_id: 'conversation-num-1',
+    },
+    streamMode: 'values',
+    version: 'v2' as const,
+  };
+
+  beforeEach(async () => {
+    conversationHistory = [];
+    const { contentParts: parts, aggregateContent: ac } =
+      createContentAggregator();
+    aggregateContent = ac;
+    contentParts = parts as t.MessageContentComplex[];
+    // Spy on the static method instead of the instance method
+    handleToolCallErrorSpy = jest.spyOn(
+      StandardGraph,
+      'handleToolCallErrorStatic'
+    );
+  });
+
+  afterEach(() => {
+    handleToolCallErrorSpy.mockRestore();
+  });
+
+  const onMessageDeltaSpy = jest.fn();
+  const onRunStepSpy = jest.fn();
+  const onRunStepCompletedSpy = jest.fn();
+
+  afterAll(() => {
+    onMessageDeltaSpy.mockReset();
+    onRunStepSpy.mockReset();
+    onRunStepCompletedSpy.mockReset();
+  });
+
+  const setupCustomHandlers = (): Record<
+    string | GraphEvents,
+    t.EventHandler
+  > => ({
+    [GraphEvents.TOOL_END]: new ToolEndHandler(),
+    [GraphEvents.CHAT_MODEL_END]: new ModelEndHandler(),
+    [GraphEvents.CHAT_MODEL_STREAM]: new ChatModelStreamHandler(),
+    [GraphEvents.ON_RUN_STEP_COMPLETED]: {
+      handle: (
+        event: GraphEvents.ON_RUN_STEP_COMPLETED,
+        data: t.StreamEventData
+      ): void => {
+        if ((data.result as t.MessageContentComplex)['type'] === 'tool_call') {
+          run.Graph?.overrideTestModel(
+            ['Looks like there was an error calling the tool.'],
+            5
+          );
+        }
+        onRunStepCompletedSpy(event, data);
+        aggregateContent({
+          event,
+          data: data as unknown as { result: t.ToolEndEvent },
+        });
+      },
+    },
+    [GraphEvents.ON_RUN_STEP]: {
+      handle: (
+        event: GraphEvents.ON_RUN_STEP,
+        data: t.StreamEventData,
+        metadata,
+        graph
+      ): void => {
+        const runStepData = data as t.RunStep;
+        onRunStepSpy(event, runStepData, metadata, graph);
+        aggregateContent({ event, data: runStepData });
+      },
+    },
+    [GraphEvents.ON_RUN_STEP_DELTA]: {
+      handle: (
+        event: GraphEvents.ON_RUN_STEP_DELTA,
+        data: t.StreamEventData
+      ): void => {
+        aggregateContent({ event, data: data as t.RunStepDeltaEvent });
+      },
+    },
+    [GraphEvents.ON_MESSAGE_DELTA]: {
+      handle: (
+        event: GraphEvents.ON_MESSAGE_DELTA,
+        data: t.StreamEventData,
+        metadata,
+        graph
+      ): void => {
+        onMessageDeltaSpy(event, data, metadata, graph);
+        aggregateContent({ event, data: data as t.MessageDeltaEvent });
+      },
+    },
+  });
+
+  test('should handle tool call errors correctly', async () => {
+    const { userName, location } = await getArgs();
+    const llmConfig = getLLMConfig(Providers.OPENAI);
+    const customHandlers = setupCustomHandlers();
+
+    // Create the run instance
+    run = await Run.create<t.IState>({
+      runId: 'test-run-id',
+      graphConfig: {
+        type: 'standard',
+        llmConfig,
+        tools: [errorTool],
+        instructions: 'You are a helpful AI assistant.',
+        additional_instructions: `The user's name is ${userName} and they are located in ${location}.`,
+      },
+      returnContent: true,
+      customHandlers,
+    });
+
+    const toolCalls: ToolCall[] = [
+      {
+        name: 'errorTool',
+        args: {
+          input: 'test input',
+        },
+        id: 'call_test123',
+        type: 'tool_call',
+      },
+    ];
+
+    const firstResponse = 'Let me try calling the tool';
+    run.Graph?.overrideTestModel([firstResponse], 5, toolCalls);
+
+    const userMessage = 'Use the error tool';
+    conversationHistory.push(new HumanMessage(userMessage));
+
+    const inputs = {
+      messages: conversationHistory,
+    };
+
+    await run.processStream(inputs, config);
+
+    // Verify handleToolCallError was called
+    expect(handleToolCallErrorSpy).toHaveBeenCalled();
+
+    // Find the tool call content part
+    const toolCallPart = contentParts.find(
+      (part) => part.type === 'tool_call'
+    ) as t.ToolCallContent | undefined;
+
+    // Verify the error message in contentParts
+    expect(toolCallPart).toBeDefined();
+    expect(toolCallPart?.tool_call?.args).toEqual(
+      JSON.stringify(toolCalls[0].args)
+    );
+    expect(toolCallPart?.tool_call?.output).toContain('Error processing tool');
+    expect(toolCallPart?.tool_call?.output).toContain(
+      'this is a test error I threw on purpose'
+    );
+  });
+});