@librechat/agents 2.2.2 → 2.2.3

package/src/graphs/Graph.ts

@@ -8,12 +8,13 @@ import { ChatOpenAI, AzureChatOpenAI } from '@langchain/openai';
 import { Runnable, RunnableConfig } from '@langchain/core/runnables';
 import { dispatchCustomEvent } from '@langchain/core/callbacks/dispatch';
 import { AIMessageChunk, ToolMessage, SystemMessage } from '@langchain/core/messages';
-import type { BaseMessage, BaseMessageFields } from '@langchain/core/messages';
+import type { BaseMessage, BaseMessageFields, UsageMetadata } from '@langchain/core/messages';
 import type * as t from '@/types';
 import { Providers, GraphEvents, GraphNodeKeys, StepTypes, Callback, ContentTypes } from '@/common';
 import { getChatModelClass, manualToolStreamProviders } from '@/llm/providers';
 import { ToolNode as CustomToolNode, toolsCondition } from '@/tools/ToolNode';
 import {
+  createPruneMessages,
   modifyDeltaProperties,
   formatArtifactPayload,
   convertMessagesToContent,
@@ -74,8 +75,13 @@ export abstract class Graph<
   stepKeyIds: Map<string, string[]> = new Map<string, string[]>();
   contentIndexMap: Map<string, number> = new Map();
   toolCallStepIds: Map<string, string> = new Map();
+  currentUsage: Partial<UsageMetadata> | undefined;
+  indexTokenCountMap: Record<string, number> = {};
+  maxContextTokens: number | undefined;
+  pruneMessages?: ReturnType<typeof createPruneMessages>;
   /** The amount of time that should pass before another consecutive API call */
   streamBuffer: number | undefined;
+  tokenCounter?: t.TokenCounter;
   signal?: AbortSignal;
 }
 
@@ -166,6 +172,10 @@ export class StandardGraph extends Graph<
     this.currentTokenType = resetIfNotEmpty(this.currentTokenType, ContentTypes.TEXT);
     this.lastToken = resetIfNotEmpty(this.lastToken, undefined);
     this.tokenTypeSwitch = resetIfNotEmpty(this.tokenTypeSwitch, undefined);
+    this.indexTokenCountMap = resetIfNotEmpty(this.indexTokenCountMap, {});
+    this.currentUsage = resetIfNotEmpty(this.currentUsage, undefined);
+    this.tokenCounter = resetIfNotEmpty(this.tokenCounter, undefined);
+    this.maxContextTokens = resetIfNotEmpty(this.maxContextTokens, undefined);
   }
 
   /* Run Step Processing */
@@ -326,6 +336,12 @@ export class StandardGraph extends Graph<
     return new ChatModelClass(options);
   }
 
+  storeUsageMetadata(finalMessage?: BaseMessage): void {
+    if (finalMessage && 'usage_metadata' in finalMessage && finalMessage.usage_metadata) {
+      this.currentUsage = finalMessage.usage_metadata as Partial<UsageMetadata>;
+    }
+  }
+
   createCallModel() {
     return async (state: t.BaseGraphState, config?: RunnableConfig): Promise<Partial<t.BaseGraphState>> => {
       const { provider = '' } = (config?.configurable as t.GraphConfig | undefined) ?? {} ;
@@ -338,9 +354,27 @@ export class StandardGraph extends Graph<
       this.config = config;
       const { messages } = state;
 
-      const finalMessages = messages;
-      const lastMessageX = finalMessages[finalMessages.length - 2];
-      const lastMessageY = finalMessages[finalMessages.length - 1];
+      let messagesToUse = messages;
+      if (!this.pruneMessages && this.tokenCounter && this.maxContextTokens && this.indexTokenCountMap[0] != null) {
+        this.pruneMessages = createPruneMessages({
+          indexTokenCountMap: this.indexTokenCountMap,
+          maxTokens: this.maxContextTokens,
+          tokenCounter: this.tokenCounter,
+          startIndex: this.startIndex,
+        });
+      }
+      if (this.pruneMessages) {
+        const { context, indexTokenCountMap } = this.pruneMessages({
+          messages,
+          usageMetadata: this.currentUsage,
+        });
+        this.indexTokenCountMap = indexTokenCountMap;
+        messagesToUse = context;
+      }
+
+      const finalMessages = messagesToUse;
+      const lastMessageX = finalMessages.length >= 2 ? finalMessages[finalMessages.length - 2] : null;
+      const lastMessageY = finalMessages.length >= 1 ? finalMessages[finalMessages.length - 1] : null;
 
       if (
         provider === Providers.BEDROCK
@@ -372,6 +406,7 @@ export class StandardGraph extends Graph<
 
       this.lastStreamCall = Date.now();
 
+      let result: Partial<t.BaseGraphState>;
       if ((this.tools?.length ?? 0) > 0 && manualToolStreamProviders.has(provider)) {
         const stream = await this.boundModel.stream(finalMessages, config);
         let finalChunk: AIMessageChunk | undefined;
@@ -385,19 +420,22 @@ export class StandardGraph extends Graph<
         }
 
         finalChunk = modifyDeltaProperties(this.provider, finalChunk);
-        return { messages: [finalChunk as AIMessageChunk] };
-      }
-
-      const finalMessage = (await this.boundModel.invoke(finalMessages, config)) as AIMessageChunk;
-      if ((finalMessage.tool_calls?.length ?? 0) > 0) {
-        finalMessage.tool_calls = finalMessage.tool_calls?.filter((tool_call) => {
-          if (!tool_call.name) {
-            return false;
-          }
-          return true;
-        });
+        result = { messages: [finalChunk as AIMessageChunk] };
+      } else {
+        const finalMessage = (await this.boundModel.invoke(finalMessages, config)) as AIMessageChunk;
+        if ((finalMessage.tool_calls?.length ?? 0) > 0) {
+          finalMessage.tool_calls = finalMessage.tool_calls?.filter((tool_call) => {
+            if (!tool_call.name) {
+              return false;
+            }
+            return true;
+          });
+        }
+        result = { messages: [finalMessage] };
       }
-      return { messages: [finalMessage] };
+      
+      this.storeUsageMetadata(result?.messages?.[0]);
+      return result;
     };
   }
 

package/src/messages/format.ts

@@ -431,3 +431,30 @@ export const formatContentStrings = (payload: Array<BaseMessage>): Array<BaseMes
 
   return result;
 };
+
+/**
+ * Adds a value at key 0 for system messages and shifts all key indices by one in an indexTokenCountMap.
+ * This is useful when adding a system message at the beginning of a conversation.
+ * 
+ * @param indexTokenCountMap - The original map of message indices to token counts
+ * @param instructionsTokenCount - The token count for the system message to add at index 0
+ * @returns A new map with the system message at index 0 and all other indices shifted by 1
+ */
+export function shiftIndexTokenCountMap(
+  indexTokenCountMap: Record<number, number>,
+  instructionsTokenCount: number
+): Record<number, number> {
+  // Create a new map to avoid modifying the original
+  const shiftedMap: Record<number, number> = {};
+  
+  // Add the system message token count at index 0
+  shiftedMap[0] = instructionsTokenCount;
+  
+  // Shift all existing indices by 1
+  for (const [indexStr, tokenCount] of Object.entries(indexTokenCountMap)) {
+    const index = Number(indexStr);
+    shiftedMap[index + 1] = tokenCount;
+  }
+  
+  return shiftedMap;
+}

package/src/messages/index.ts

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

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);