@librechat/agents 2.2.1 → 2.2.3

package/src/messages/formatMessage.test.ts

@@ -0,0 +1,277 @@
+import { HumanMessage, AIMessage, SystemMessage } from '@langchain/core/messages';
+import { formatMessage, formatLangChainMessages, formatFromLangChain } from './format';
+
+const NO_PARENT = '00000000-0000-0000-0000-000000000000';
+
+describe('formatMessage', () => {
+  it('formats user message', () => {
+    const input = {
+      message: {
+        sender: 'user',
+        text: 'Hello',
+      },
+      userName: 'John',
+    };
+    const result = formatMessage(input);
+    expect(result).toEqual({
+      role: 'user',
+      content: 'Hello',
+      name: 'John',
+    });
+  });
+
+  it('sanitizes the name by replacing invalid characters (per OpenAI)', () => {
+    const input = {
+      message: {
+        sender: 'user',
+        text: 'Hello',
+      },
+      userName: ' John$Doe@Example! ',
+    };
+    const result = formatMessage(input);
+    expect(result).toEqual({
+      role: 'user',
+      content: 'Hello',
+      name: '_John_Doe_Example__',
+    });
+  });
+
+  it('trims the name to a maximum length of 64 characters', () => {
+    const longName = 'a'.repeat(100);
+    const input = {
+      message: {
+        sender: 'user',
+        text: 'Hello',
+      },
+      userName: longName,
+    };
+    const result = formatMessage(input);
+    expect(result.name?.length).toBe(64);
+    expect(result.name).toBe('a'.repeat(64));
+  });
+
+  it('formats a realistic user message', () => {
+    const input = {
+      message: {
+        _id: '6512cdfb92cbf69fea615331',
+        messageId: 'b620bf73-c5c3-4a38-b724-76886aac24c4',
+        __v: 0,
+        conversationId: '5c23d24f-941f-4aab-85df-127b596c8aa5',
+        createdAt: Date.now(),
+        error: false,
+        finish_reason: null,
+        isCreatedByUser: true,
+        model: null,
+        parentMessageId: NO_PARENT,
+        sender: 'User',
+        text: 'hi',
+        tokenCount: 5,
+        unfinished: false,
+        updatedAt: Date.now(),
+        user: '6512cdf475f05c86d44c31d2',
+      },
+      userName: 'John',
+    };
+    const result = formatMessage(input);
+    expect(result).toEqual({
+      role: 'user',
+      content: 'hi',
+      name: 'John',
+    });
+  });
+
+  it('formats assistant message', () => {
+    const input = {
+      message: {
+        sender: 'assistant',
+        text: 'Hi there',
+      },
+      assistantName: 'Assistant',
+    };
+    const result = formatMessage(input);
+    expect(result).toEqual({
+      role: 'assistant',
+      content: 'Hi there',
+      name: 'Assistant',
+    });
+  });
+
+  it('formats system message', () => {
+    const input = {
+      message: {
+        role: 'system',
+        text: 'Hi there',
+      },
+    };
+    const result = formatMessage(input);
+    expect(result).toEqual({
+      role: 'system',
+      content: 'Hi there',
+    });
+  });
+
+  it('formats user message with langChain', () => {
+    const input = {
+      message: {
+        sender: 'user',
+        text: 'Hello',
+      },
+      userName: 'John',
+      langChain: true,
+    };
+    const result = formatMessage(input);
+    expect(result).toBeInstanceOf(HumanMessage);
+    expect(result.lc_kwargs.content).toEqual(input.message.text);
+    expect(result.lc_kwargs.name).toEqual(input.userName);
+  });
+
+  it('formats assistant message with langChain', () => {
+    const input = {
+      message: {
+        sender: 'assistant',
+        text: 'Hi there',
+      },
+      assistantName: 'Assistant',
+      langChain: true,
+    };
+    const result = formatMessage(input);
+    expect(result).toBeInstanceOf(AIMessage);
+    expect(result.lc_kwargs.content).toEqual(input.message.text);
+    expect(result.lc_kwargs.name).toEqual(input.assistantName);
+  });
+
+  it('formats system message with langChain', () => {
+    const input = {
+      message: {
+        role: 'system',
+        text: 'This is a system message.',
+      },
+      langChain: true,
+    };
+    const result = formatMessage(input);
+    expect(result).toBeInstanceOf(SystemMessage);
+    expect(result.lc_kwargs.content).toEqual(input.message.text);
+  });
+
+  it('formats langChain messages into OpenAI payload format', () => {
+    const human = {
+      message: new HumanMessage({
+        content: 'Hello',
+      }),
+    };
+    const system = {
+      message: new SystemMessage({
+        content: 'Hello',
+      }),
+    };
+    const ai = {
+      message: new AIMessage({
+        content: 'Hello',
+      }),
+    };
+    const humanResult = formatMessage(human);
+    const systemResult = formatMessage(system);
+    const aiResult = formatMessage(ai);
+    expect(humanResult).toEqual({
+      role: 'user',
+      content: 'Hello',
+    });
+    expect(systemResult).toEqual({
+      role: 'system',
+      content: 'Hello',
+    });
+    expect(aiResult).toEqual({
+      role: 'assistant',
+      content: 'Hello',
+    });
+  });
+});
+
+describe('formatLangChainMessages', () => {
+  it('formats an array of messages for LangChain', () => {
+    const messages = [
+      {
+        role: 'system',
+        content: 'This is a system message',
+      },
+      {
+        sender: 'user',
+        text: 'Hello',
+      },
+      {
+        sender: 'assistant',
+        text: 'Hi there',
+      },
+    ];
+    const formatOptions = {
+      userName: 'John',
+      assistantName: 'Assistant',
+    };
+    const result = formatLangChainMessages(messages, formatOptions);
+    expect(result).toHaveLength(3);
+    expect(result[0]).toBeInstanceOf(SystemMessage);
+    expect(result[1]).toBeInstanceOf(HumanMessage);
+    expect(result[2]).toBeInstanceOf(AIMessage);
+
+    expect(result[0].lc_kwargs.content).toEqual(messages[0].content);
+    expect(result[1].lc_kwargs.content).toEqual(messages[1].text);
+    expect(result[2].lc_kwargs.content).toEqual(messages[2].text);
+
+    expect(result[1].lc_kwargs.name).toEqual(formatOptions.userName);
+    expect(result[2].lc_kwargs.name).toEqual(formatOptions.assistantName);
+  });
+
+  describe('formatFromLangChain', () => {
+    it('should merge kwargs and additional_kwargs', () => {
+      const message = {
+        kwargs: {
+          content: 'some content',
+          name: 'dan',
+          additional_kwargs: {
+            function_call: {
+              name: 'dall-e',
+              arguments: '{\n  "input": "Subject: hedgehog, Style: cute"\n}',
+            },
+          },
+        },
+      };
+
+      const expected = {
+        content: 'some content',
+        name: 'dan',
+        function_call: {
+          name: 'dall-e',
+          arguments: '{\n  "input": "Subject: hedgehog, Style: cute"\n}',
+        },
+      };
+
+      expect(formatFromLangChain(message)).toEqual(expected);
+    });
+
+    it('should handle messages without additional_kwargs', () => {
+      const message = {
+        kwargs: {
+          content: 'some content',
+          name: 'dan',
+        },
+      };
+
+      const expected = {
+        content: 'some content',
+        name: 'dan',
+      };
+
+      expect(formatFromLangChain(message)).toEqual(expected);
+    });
+
+    it('should handle empty messages', () => {
+      const message = {
+        kwargs: {},
+      };
+
+      const expected = {};
+
+      expect(formatFromLangChain(message)).toEqual(expected);
+    });
+  });
+});

package/src/messages/index.ts

@@ -0,0 +1,3 @@
+export * from './core';
+export * from './prune';
+export * from './format';

package/src/messages/prune.ts

@@ -0,0 +1,167 @@
+import type { BaseMessage, UsageMetadata } from '@langchain/core/messages';
+import type { TokenCounter } from '@/types/run';
+export type PruneMessagesFactoryParams = {
+  maxTokens: number;
+  startIndex: number;
+  tokenCounter: TokenCounter;
+  indexTokenCountMap: Record<string, number>;
+};
+export type PruneMessagesParams = {
+  messages: BaseMessage[];
+  usageMetadata?: Partial<UsageMetadata>;
+}
+
+/**
+ * Calculates the total tokens from a single usage object
+ * 
+ * @param usage The usage metadata object containing token information
+ * @returns An object containing the total input and output tokens
+ */
+function calculateTotalTokens(usage: Partial<UsageMetadata>): UsageMetadata {
+  const baseInputTokens = Number(usage.input_tokens) || 0;
+  const cacheCreation = Number(usage.input_token_details?.cache_creation) || 0;
+  const cacheRead = Number(usage.input_token_details?.cache_read) || 0;
+  
+  const totalInputTokens = baseInputTokens + cacheCreation + cacheRead;
+  const totalOutputTokens = Number(usage.output_tokens) || 0;
+
+  return {
+    input_tokens: totalInputTokens,
+    output_tokens: totalOutputTokens,
+    total_tokens: totalInputTokens + totalOutputTokens
+  };
+}
+
+/**
+ * Processes an array of messages and returns a context of messages that fit within a specified token limit.
+ * It iterates over the messages from newest to oldest, adding them to the context until the token limit is reached.
+ * 
+ * @param options Configuration options for processing messages
+ * @returns Object containing the message context, remaining tokens, messages not included, and summary index
+ */
+function getMessagesWithinTokenLimit({
+  messages: _messages,
+  maxContextTokens,
+  indexTokenCountMap,
+}: {
+  messages: BaseMessage[];
+  maxContextTokens: number;
+  indexTokenCountMap: Record<string, number>;
+}): {
+  context: BaseMessage[];
+  remainingContextTokens: number;
+  messagesToRefine: BaseMessage[];
+  summaryIndex: number;
+} {
+  // Every reply is primed with <|start|>assistant<|message|>, so we
+  // start with 3 tokens for the label after all messages have been counted.
+  let summaryIndex = -1;
+  let currentTokenCount = 3;
+  const instructions = _messages?.[0]?.getType() === 'system' ? _messages[0] : undefined;
+  const instructionsTokenCount = instructions != null ? indexTokenCountMap[0] : 0;
+  let remainingContextTokens = maxContextTokens - instructionsTokenCount;
+  const messages = [..._messages];
+  const context: BaseMessage[] = [];
+
+  if (currentTokenCount < remainingContextTokens) {
+    let currentIndex = messages.length;
+    while (messages.length > 0 && currentTokenCount < remainingContextTokens && currentIndex > 1) {
+      currentIndex--;
+      if (messages.length === 1 && instructions) {
+        break;
+      }
+      const poppedMessage = messages.pop();
+      if (!poppedMessage) continue;
+      
+      const tokenCount = indexTokenCountMap[currentIndex] || 0;
+
+      if ((currentTokenCount + tokenCount) <= remainingContextTokens) {
+        context.push(poppedMessage);
+        currentTokenCount += tokenCount;
+      } else {
+        messages.push(poppedMessage);
+        break;
+      }
+    }
+  }
+
+  if (instructions && _messages.length > 0) {
+    context.push(_messages[0] as BaseMessage);
+    messages.shift();
+  }
+
+  const prunedMemory = messages;
+  summaryIndex = prunedMemory.length - 1;
+  remainingContextTokens -= currentTokenCount;
+
+  return {
+    summaryIndex,
+    remainingContextTokens,
+    context: context.reverse(),
+    messagesToRefine: prunedMemory,
+  };
+}
+
+function checkValidNumber(value: unknown): value is number {
+  return typeof value === 'number' && !isNaN(value) && value > 0;
+}
+
+export function createPruneMessages(factoryParams: PruneMessagesFactoryParams) {
+  const indexTokenCountMap = { ...factoryParams.indexTokenCountMap };
+  let lastTurnStartIndex = factoryParams.startIndex;
+  let totalTokens = (Object.values(indexTokenCountMap)).reduce((a, b) => a + b, 0);
+  return function pruneMessages(params: PruneMessagesParams): {
+    context: BaseMessage[];
+    indexTokenCountMap: Record<string, number>;
+  } {
+    let currentUsage: UsageMetadata | undefined;
+    if (params.usageMetadata && (
+      checkValidNumber(params.usageMetadata.input_tokens)
+      || (
+        checkValidNumber(params.usageMetadata.input_token_details)
+        && (
+          checkValidNumber(params.usageMetadata.input_token_details.cache_creation)
+          || checkValidNumber(params.usageMetadata.input_token_details.cache_read)
+        )
+      )
+    ) && checkValidNumber(params.usageMetadata.output_tokens)) {
+      currentUsage = calculateTotalTokens(params.usageMetadata);
+      totalTokens = currentUsage.total_tokens;
+    }
+
+    for (let i = lastTurnStartIndex; i < params.messages.length; i++) {
+      const message = params.messages[i];
+      if (i === lastTurnStartIndex && indexTokenCountMap[i] === undefined && currentUsage) {
+        indexTokenCountMap[i] = currentUsage.output_tokens;
+      } else if (indexTokenCountMap[i] === undefined) {
+        indexTokenCountMap[i] = factoryParams.tokenCounter(message);
+        totalTokens += indexTokenCountMap[i];
+      }
+    }
+
+    // If `currentUsage` is defined, we need to distribute the current total tokensto our `indexTokenCountMap`,
+    // for all message index keys before `lastTurnStartIndex`, as it has the most accurate count for those messages.
+    // We must distribute it in a weighted manner, so that the total token count is equal to `currentUsage.total_tokens`,
+    // relative the manually counted tokens in `indexTokenCountMap`.
+    if (currentUsage) {
+      const totalIndexTokens = Object.values(indexTokenCountMap).reduce((a, b) => a + b, 0);
+      const ratio = currentUsage.total_tokens / totalIndexTokens;
+      for (const key in indexTokenCountMap) {
+        indexTokenCountMap[key] = Math.round(indexTokenCountMap[key] * ratio);
+      }
+    }
+
+    lastTurnStartIndex = params.messages.length;
+    if (totalTokens <= factoryParams.maxTokens) {
+      return { context: params.messages, indexTokenCountMap };
+    }
+
+    const { context } = getMessagesWithinTokenLimit({
+      maxContextTokens: factoryParams.maxTokens,
+      messages: params.messages,
+      indexTokenCountMap,
+    });
+
+    return { context, indexTokenCountMap };
+  }
+}

package/src/messages/shiftIndexTokenCountMap.test.ts

@@ -0,0 +1,81 @@
+import { shiftIndexTokenCountMap } from './format';
+
+describe('shiftIndexTokenCountMap', () => {
+  it('should add a system message token count at index 0 and shift all other indices', () => {
+    const originalMap: Record<number, number> = {
+      0: 10,
+      1: 20,
+      2: 30
+    };
+    
+    const systemMessageTokenCount = 15;
+    
+    const result = shiftIndexTokenCountMap(originalMap, systemMessageTokenCount);
+    
+    // Check that the system message token count is at index 0
+    expect(result[0]).toBe(15);
+    
+    // Check that all other indices are shifted by 1
+    expect(result[1]).toBe(10);
+    expect(result[2]).toBe(20);
+    expect(result[3]).toBe(30);
+    
+    // Check that the original map is not modified
+    expect(originalMap[0]).toBe(10);
+    expect(originalMap[1]).toBe(20);
+    expect(originalMap[2]).toBe(30);
+  });
+  
+  it('should handle an empty map', () => {
+    const emptyMap: Record<number, number> = {};
+    const systemMessageTokenCount = 15;
+    
+    const result = shiftIndexTokenCountMap(emptyMap, systemMessageTokenCount);
+    
+    // Check that only the system message token count is in the result
+    expect(Object.keys(result).length).toBe(1);
+    expect(result[0]).toBe(15);
+  });
+  
+  it('should handle non-sequential indices', () => {
+    const nonSequentialMap: Record<number, number> = {
+      0: 10,
+      2: 20,
+      5: 30
+    };
+    
+    const systemMessageTokenCount = 15;
+    
+    const result = shiftIndexTokenCountMap(nonSequentialMap, systemMessageTokenCount);
+    
+    // Check that the system message token count is at index 0
+    expect(result[0]).toBe(15);
+    
+    // Check that all other indices are shifted by 1
+    expect(result[1]).toBe(10);
+    expect(result[3]).toBe(20);
+    expect(result[6]).toBe(30);
+  });
+  
+  it('should handle string keys', () => {
+    // TypeScript will convert string keys to numbers when accessing the object
+    const mapWithStringKeys: Record<string, number> = {
+      '0': 10,
+      '1': 20,
+      '2': 30
+    };
+    
+    const systemMessageTokenCount = 15;
+    
+    // Cast to Record<number, number> to match the function signature
+    const result = shiftIndexTokenCountMap(mapWithStringKeys as unknown as Record<number, number>, systemMessageTokenCount);
+    
+    // Check that the system message token count is at index 0
+    expect(result[0]).toBe(15);
+    
+    // Check that all other indices are shifted by 1
+    expect(result[1]).toBe(10);
+    expect(result[2]).toBe(20);
+    expect(result[3]).toBe(30);
+  });
+});

package/src/run.ts

@@ -1,13 +1,17 @@
 // src/run.ts
+import { zodToJsonSchema } from "zod-to-json-schema";
 import { PromptTemplate } from '@langchain/core/prompts';
 import { AzureChatOpenAI, ChatOpenAI } from '@langchain/openai';
+import { SystemMessage } from '@langchain/core/messages';
 import type { BaseMessage, MessageContentComplex } from '@langchain/core/messages';
 import type { ClientCallbacks, SystemCallbacks } from '@/graphs/Graph';
 import type { RunnableConfig } from '@langchain/core/runnables';
 import type * as t from '@/types';
 import { GraphEvents, Providers, Callback } from '@/common';
 import { manualToolStreamProviders } from '@/llm/providers';
+import { shiftIndexTokenCountMap } from '@/messages/format';
 import { createTitleRunnable } from '@/utils/title';
+import { createTokenCounter } from '@/utils/tokens';
 import { StandardGraph } from '@/graphs/Graph';
 import { HandlerRegistry } from '@/events';
 import { isOpenAILike } from '@/utils/llm';
@@ -106,6 +110,28 @@ export class Run<T extends t.BaseGraphState> {
       throw new Error('Run ID not provided');
     }
 
+    const tokenCounter = streamOptions?.tokenCounter ?? (streamOptions?.indexTokenCountMap ? await createTokenCounter() : undefined);
+    const toolTokens = tokenCounter ? (this.Graph.tools?.reduce((acc, tool) => {
+      if (!tool.schema) {
+        return acc;
+      }
+  
+      const jsonSchema = zodToJsonSchema(tool.schema.describe(tool.description ?? ''), tool.name);
+      return acc + tokenCounter(new SystemMessage(JSON.stringify(jsonSchema)));
+    }, 0) ?? 0) : 0;
+    let instructionTokens = toolTokens;
+    if (this.Graph.systemMessage && tokenCounter) {
+      instructionTokens += tokenCounter(this.Graph.systemMessage);
+    }
+    if (instructionTokens > 0) {
+      this.Graph.indexTokenCountMap = shiftIndexTokenCountMap(streamOptions?.indexTokenCountMap ?? {}, instructionTokens);
+    } else {
+      this.Graph.indexTokenCountMap = streamOptions?.indexTokenCountMap ?? {};
+    }
+
+    this.Graph.maxContextTokens = streamOptions?.maxContextTokens;
+    this.Graph.tokenCounter = tokenCounter;
+
     config.run_id = this.id;
     config.configurable = Object.assign(config.configurable ?? {}, { run_id: this.id, provider: this.provider });
 

package/src/scripts/code_exec_simple.ts

@@ -1,12 +1,13 @@
 // src/scripts/cli.ts
 import { config } from 'dotenv';
 config();
-import { HumanMessage, AIMessage, BaseMessage } from '@langchain/core/messages';
+import { HumanMessage, BaseMessage } from '@langchain/core/messages';
 import { TavilySearchResults } from '@langchain/community/tools/tavily_search';
 import type * as t from '@/types';
 import { ChatModelStreamHandler, createContentAggregator } from '@/stream';
-import { ToolEndHandler, ModelEndHandler, createMetadataAggregator } from '@/events';
 import { createCodeExecutionTool } from '@/tools/CodeExecutor';
+import { ToolEndHandler, ModelEndHandler } from '@/events';
+import { createTokenCounter } from '@/utils/tokens';
 import { getLLMConfig } from '@/utils/llmConfig';
 import { getArgs } from '@/scripts/args';
 import { GraphEvents } from '@/common';
@@ -58,19 +59,22 @@ async function testCodeExecution(): Promise<void> {
   };
 
   const llmConfig = getLLMConfig(provider);
+  const instructions = 'You are a friendly AI assistant with coding capabilities. Always address the user by their name.';
+  const additional_instructions = `The user's name is ${userName} and they are located in ${location}.`;
 
-  const run = await Run.create<t.IState>({
+  const runConfig: t.RunConfig = {
     runId: 'message-num-1',
     graphConfig: {
       type: 'standard',
       llmConfig,
       tools: [new TavilySearchResults(), createCodeExecutionTool()],
-      instructions: 'You are a friendly AI assistant with coding capabilities. Always address the user by their name.',
-      additional_instructions: `The user's name is ${userName} and they are located in ${location}.`,
+      instructions,
+      additional_instructions,
     },
     returnContent: true,
     customHandlers,
-  });
+  };
+  const run = await Run.create<t.IState>(runConfig);
 
   const config = {
     configurable: {
@@ -86,13 +90,22 @@ async function testCodeExecution(): Promise<void> {
   // const userMessage1 = `how much memory is this (its in bytes) in MB? 31192000`;
   // const userMessage1 = `can you show me a good use case for rscript by running some code`;
   const userMessage1 = `Run hello world in french and in english, using python. please run 2 parallel code executions.`;
+  const humanMessage = new HumanMessage(userMessage1);
+  const tokenCounter = await createTokenCounter();
+  const indexTokenCountMap = {
+    0: tokenCounter(humanMessage),
+  };
 
-  conversationHistory.push(new HumanMessage(userMessage1));
+  conversationHistory.push(humanMessage);
 
   let inputs = {
     messages: conversationHistory,
   };
-  const finalContentParts1 = await run.processStream(inputs, config);
+  const finalContentParts1 = await run.processStream(inputs, config, {
+    maxContextTokens: 8000,
+    indexTokenCountMap,
+    tokenCounter,
+  });
   const finalMessages1 = run.getRunMessages();
   if (finalMessages1) {
     conversationHistory.push(...finalMessages1);