@librechat/agents 2.3.7 → 2.3.9

package/src/messages/prune.ts

@@ -1,8 +1,10 @@
 import { AIMessage, BaseMessage, UsageMetadata } from '@langchain/core/messages';
-import type { ThinkingContentText, MessageContentComplex } from '@/types/stream';
+import type { ThinkingContentText, MessageContentComplex, ReasoningContentText } from '@/types/stream';
 import type { TokenCounter } from '@/types/run';
-import { ContentTypes } from '@/common';
+import { ContentTypes, Providers } from '@/common';
+
 export type PruneMessagesFactoryParams = {
+  provider?: Providers;
   maxTokens: number;
   startIndex: number;
   tokenCounter: TokenCounter;
@@ -20,7 +22,7 @@ function isIndexInContext(arrayA: unknown[], arrayB: unknown[], targetIndex: num
   return targetIndex >= startingIndexInA;
 }
 
-function addThinkingBlock(message: AIMessage, thinkingBlock: ThinkingContentText): MessageContentComplex[] {
+function addThinkingBlock(message: AIMessage, thinkingBlock: ThinkingContentText | ReasoningContentText): MessageContentComplex[] {
   const content: MessageContentComplex[] = Array.isArray(message.content)
     ? message.content as MessageContentComplex[]
     : [{
@@ -52,6 +54,13 @@ export function calculateTotalTokens(usage: Partial<UsageMetadata>): UsageMetada
   };
 }
 
+export type PruningResult = {
+  context: BaseMessage[];
+  remainingContextTokens: number;
+  messagesToRefine: BaseMessage[];
+  thinkingStartIndex?: number;
+};
+
 /**
  * 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.
@@ -65,20 +74,19 @@ export function getMessagesWithinTokenLimit({
   indexTokenCountMap,
   startType: _startType,
   thinkingEnabled,
-  /** We may need to use this when recalculating */
   tokenCounter,
+  thinkingStartIndex: _thinkingStartIndex = -1,
+  reasoningType = ContentTypes.THINKING,
 }: {
   messages: BaseMessage[];
   maxContextTokens: number;
   indexTokenCountMap: Record<string, number | undefined>;
-  tokenCounter: TokenCounter;
-  startType?: string;
+  startType?: string | string[];
   thinkingEnabled?: boolean;
-}): {
-  context: BaseMessage[];
-  remainingContextTokens: number;
-  messagesToRefine: BaseMessage[];
-} {
+  tokenCounter: TokenCounter;
+  thinkingStartIndex?: number;
+  reasoningType?: ContentTypes.THINKING | ContentTypes.REASONING_CONTENT;
+}): PruningResult {
   // Every reply is primed with <|start|>assistant<|message|>, so we
   // start with 3 tokens for the label after all messages have been counted.
   let currentTokenCount = 3;
@@ -96,12 +104,19 @@ export function getMessagesWithinTokenLimit({
    * */
   let context: Array<BaseMessage | undefined> = [];
 
-  let thinkingStartIndex = -1;
+  let thinkingStartIndex = _thinkingStartIndex;
   let thinkingEndIndex = -1;
-  let thinkingBlock: ThinkingContentText | undefined;
+  let thinkingBlock: ThinkingContentText | ReasoningContentText | undefined;
   const endIndex = instructions != null ? 1 : 0;
   const prunedMemory: BaseMessage[] = [];
 
+  if (_thinkingStartIndex > -1) {
+    const thinkingMessageContent = _messages[_thinkingStartIndex]?.content;
+    if (Array.isArray(thinkingMessageContent)) {
+      thinkingBlock = thinkingMessageContent.find((content) => content.type === reasoningType) as ThinkingContentText | undefined;
+    }
+  }
+
   if (currentTokenCount < remainingContextTokens) {
     let currentIndex = messages.length;
     while (messages.length > 0 && currentTokenCount < remainingContextTokens && currentIndex > endIndex) {
@@ -116,7 +131,7 @@ export function getMessagesWithinTokenLimit({
         thinkingEndIndex = currentIndex;
       }
       if (thinkingEndIndex > -1 && !thinkingBlock  && thinkingStartIndex < 0 && messageType === 'ai' && Array.isArray(poppedMessage.content)) {
-        thinkingBlock = (poppedMessage.content.find((content) => content.type === ContentTypes.THINKING)) as ThinkingContentText | undefined;
+        thinkingBlock = (poppedMessage.content.find((content) => content.type === reasoningType)) as ThinkingContentText | undefined;
         thinkingStartIndex = thinkingBlock != null ? currentIndex : -1;
       }
       /** False start, the latest message was not part of a multi-assistant/tool sequence of messages */
@@ -135,22 +150,34 @@ export function getMessagesWithinTokenLimit({
         currentTokenCount += tokenCount;
       } else {
         prunedMemory.push(poppedMessage);
-        if (thinkingEndIndex > -1) {
+        if (thinkingEndIndex > -1 && thinkingStartIndex < 0) {
           continue;
         }
         break;
       }
     }
 
-    if (thinkingEndIndex > -1 && context[context.length - 1]?.getType() === 'tool') {
-      startType = 'ai';
+    if (context[context.length - 1]?.getType() === 'tool') {
+      startType = ['ai', 'human'];
     }
 
-    if (startType != null && startType && context.length > 0) {
-      const requiredTypeIndex = context.findIndex(msg => msg?.getType() === startType);
+    if (startType != null && startType.length > 0 && context.length > 0) {
+      let requiredTypeIndex = -1;
+
+      let totalTokens = 0;
+      for (let i = context.length - 1; i >= 0; i--) {
+        const currentType = context[i]?.getType() ?? '';
+        if (Array.isArray(startType) ? startType.includes(currentType) : currentType === startType) {
+          requiredTypeIndex = i + 1;
+          break;
+        }
+        const originalIndex = originalLength - 1 - i;
+        totalTokens += indexTokenCountMap[originalIndex] ?? 0;
+      }
 
       if (requiredTypeIndex > 0) {
-        context = context.slice(requiredTypeIndex);
+        currentTokenCount -= totalTokens;
+        context = context.slice(0, requiredTypeIndex);
       }
     }
   }
@@ -161,12 +188,16 @@ export function getMessagesWithinTokenLimit({
   }
 
   remainingContextTokens -= currentTokenCount;
-  const result = {
+  const result: PruningResult = {
     remainingContextTokens,
     context: [] as BaseMessage[],
     messagesToRefine: prunedMemory,
   };
 
+  if (thinkingStartIndex > -1) {
+    result.thinkingStartIndex = thinkingStartIndex;
+  }
+
   if (prunedMemory.length === 0 || thinkingEndIndex < 0 || (thinkingStartIndex > -1 && isIndexInContext(_messages, context, thinkingStartIndex))) {
     // we reverse at this step to ensure the context is in the correct order for the model, and we need to work backwards
     result.context = context.reverse() as BaseMessage[];
@@ -270,6 +301,7 @@ export function createPruneMessages(factoryParams: PruneMessagesFactoryParams) {
   let lastTurnStartIndex = factoryParams.startIndex;
   let lastCutOffIndex = 0;
   let totalTokens = (Object.values(indexTokenCountMap)).reduce((a, b) => a + b, 0);
+  let runThinkingStartIndex = -1;
   return function pruneMessages(params: PruneMessagesParams): {
     context: BaseMessage[];
     indexTokenCountMap: Record<string, number>;
@@ -339,15 +371,19 @@ export function createPruneMessages(factoryParams: PruneMessagesFactoryParams) {
       return { context: params.messages, indexTokenCountMap };
     }
 
-    const { context } = getMessagesWithinTokenLimit({
+    const { context, thinkingStartIndex } = getMessagesWithinTokenLimit({
       maxContextTokens: factoryParams.maxTokens,
       messages: params.messages,
       indexTokenCountMap,
       startType: params.startType,
       thinkingEnabled: factoryParams.thinkingEnabled,
       tokenCounter: factoryParams.tokenCounter,
+      reasoningType: factoryParams.provider === Providers.BEDROCK ? ContentTypes.REASONING_CONTENT : ContentTypes.THINKING,
+      thinkingStartIndex: factoryParams.thinkingEnabled === true ? runThinkingStartIndex : undefined,
     });
-    lastCutOffIndex = Math.max(params.messages.length - context.length, 0);
+    runThinkingStartIndex = thinkingStartIndex ?? -1;
+    /** The index is the first value of `context`, index relative to `params.messages` */
+    lastCutOffIndex = Math.max(params.messages.length - (context.length - (context[0]?.getType() === 'system' ? 1 : 0)), 0);
 
     return { context, indexTokenCountMap };
   };

package/src/specs/prune.test.ts

@@ -1,7 +1,7 @@
 // src/specs/prune.test.ts
 import { config } from 'dotenv';
 config();
-import { HumanMessage, AIMessage, SystemMessage, BaseMessage } from '@langchain/core/messages';
+import { HumanMessage, AIMessage, SystemMessage, BaseMessage, ToolMessage } from '@langchain/core/messages';
 import type { RunnableConfig } from '@langchain/core/runnables';
 import type { UsageMetadata } from '@langchain/core/messages';
 import type * as t from '@/types';
@@ -512,6 +512,187 @@ describe('Prune Messages Tests', () => {
     });
   });
 
+  describe('Tool Message Handling', () => {
+    it('should ensure context does not start with a tool message by finding an AI message', () => {
+      const tokenCounter = createTestTokenCounter();
+      const messages = [
+        new SystemMessage('System instruction'),
+        new AIMessage('AI message 1'),
+        new ToolMessage({ content: 'Tool result 1', tool_call_id: 'tool1' }),
+        new AIMessage('AI message 2'),
+        new ToolMessage({ content: 'Tool result 2', tool_call_id: 'tool2' })
+      ];
+
+      const indexTokenCountMap = {
+        0: 17, // System instruction
+        1: 12, // AI message 1
+        2: 13, // Tool result 1
+        3: 12, // AI message 2
+        4: 13  // Tool result 2
+      };
+
+      // Create a pruneMessages function with a token limit that will only include the last few messages
+      const pruneMessages = createPruneMessages({
+        maxTokens: 58, // Only enough for system + last 3 messages + 3, but should not include a parent-less tool message
+        startIndex: 0,
+        tokenCounter,
+        indexTokenCountMap: { ...indexTokenCountMap }
+      });
+
+      const result = pruneMessages({ messages });
+
+      // The context should include the system message, AI message 2, and Tool result 2
+      // It should NOT start with Tool result 2 alone
+      expect(result.context.length).toBe(3);
+      expect(result.context[0]).toBe(messages[0]); // System message
+      expect(result.context[1]).toBe(messages[3]); // AI message 2
+      expect(result.context[2]).toBe(messages[4]); // Tool result 2
+    });
+
+    it('should ensure context does not start with a tool message by finding a human message', () => {
+      const tokenCounter = createTestTokenCounter();
+      const messages = [
+        new SystemMessage('System instruction'),
+        new HumanMessage('Human message 1'),
+        new AIMessage('AI message 1'),
+        new ToolMessage({ content: 'Tool result 1', tool_call_id: 'tool1' }),
+        new HumanMessage('Human message 2'),
+        new ToolMessage({ content: 'Tool result 2', tool_call_id: 'tool2' })
+      ];
+
+      const indexTokenCountMap = {
+        0: 17, // System instruction
+        1: 15, // Human message 1
+        2: 12, // AI message 1
+        3: 13, // Tool result 1
+        4: 15, // Human message 2
+        5: 13  // Tool result 2
+      };
+
+      // Create a pruneMessages function with a token limit that will only include the last few messages
+      const pruneMessages = createPruneMessages({
+        maxTokens: 48, // Only enough for system + last 2 messages
+        startIndex: 0,
+        tokenCounter,
+        indexTokenCountMap: { ...indexTokenCountMap }
+      });
+
+      const result = pruneMessages({ messages });
+
+      // The context should include the system message, Human message 2, and Tool result 2
+      // It should NOT start with Tool result 2 alone
+      expect(result.context.length).toBe(3);
+      expect(result.context[0]).toBe(messages[0]); // System message
+      expect(result.context[1]).toBe(messages[4]); // Human message 2
+      expect(result.context[2]).toBe(messages[5]); // Tool result 2
+    });
+
+    it('should handle the case where a tool message is followed by an AI message', () => {
+      const tokenCounter = createTestTokenCounter();
+      const messages = [
+        new SystemMessage('System instruction'),
+        new HumanMessage('Human message'),
+        new AIMessage('AI message with tool use'),
+        new ToolMessage({ content: 'Tool result', tool_call_id: 'tool1' }),
+        new AIMessage('AI message after tool')
+      ];
+
+      const indexTokenCountMap = {
+        0: 17, // System instruction
+        1: 13, // Human message
+        2: 22, // AI message with tool use
+        3: 11, // Tool result
+        4: 19  // AI message after tool
+      };
+
+      const pruneMessages = createPruneMessages({
+        maxTokens: 50,
+        startIndex: 0,
+        tokenCounter,
+        indexTokenCountMap: { ...indexTokenCountMap }
+      });
+
+      const result = pruneMessages({ messages });
+
+      expect(result.context.length).toBe(2);
+      expect(result.context[0]).toBe(messages[0]); // System message
+      expect(result.context[1]).toBe(messages[4]); // AI message after tool
+    });
+
+    it('should handle the case where a tool message is followed by a human message', () => {
+      const tokenCounter = createTestTokenCounter();
+      const messages = [
+        new SystemMessage('System instruction'),
+        new HumanMessage('Human message 1'),
+        new AIMessage('AI message with tool use'),
+        new ToolMessage({ content: 'Tool result', tool_call_id: 'tool1' }),
+        new HumanMessage('Human message 2')
+      ];
+
+      const indexTokenCountMap = {
+        0: 17, // System instruction
+        1: 15, // Human message 1
+        2: 22, // AI message with tool use
+        3: 11, // Tool result
+        4: 15  // Human message 2
+      };
+
+      const pruneMessages = createPruneMessages({
+        maxTokens: 46,
+        startIndex: 0,
+        tokenCounter,
+        indexTokenCountMap: { ...indexTokenCountMap }
+      });
+
+      const result = pruneMessages({ messages });
+
+      expect(result.context.length).toBe(2);
+      expect(result.context[0]).toBe(messages[0]); // System message
+      expect(result.context[1]).toBe(messages[4]); // Human message 2
+    });
+
+    it('should handle complex sequence with multiple tool messages', () => {
+      const tokenCounter = createTestTokenCounter();
+      const messages = [
+        new SystemMessage('System instruction'),
+        new HumanMessage('Human message 1'),
+        new AIMessage('AI message 1 with tool use'),
+        new ToolMessage({ content: 'Tool result 1', tool_call_id: 'tool1' }),
+        new AIMessage('AI message 2 with tool use'),
+        new ToolMessage({ content: 'Tool result 2', tool_call_id: 'tool2' }),
+        new AIMessage('AI message 3 with tool use'),
+        new ToolMessage({ content: 'Tool result 3', tool_call_id: 'tool3' })
+      ];
+
+      const indexTokenCountMap = {
+        0: 17, // System instruction
+        1: 15, // Human message 1
+        2: 26, // AI message 1 with tool use
+        3: 13, // Tool result 1
+        4: 26, // AI message 2 with tool use
+        5: 13, // Tool result 2
+        6: 26, // AI message 3 with tool use
+        7: 13  // Tool result 3
+      };
+
+      const pruneMessages = createPruneMessages({
+        maxTokens: 111,
+        startIndex: 0,
+        tokenCounter,
+        indexTokenCountMap: { ...indexTokenCountMap }
+      });
+
+      const result = pruneMessages({ messages });
+
+      expect(result.context.length).toBe(5);
+      expect(result.context[0]).toBe(messages[0]); // System message
+      expect(result.context[1]).toBe(messages[4]); // AI message 2 with tool use
+      expect(result.context[2]).toBe(messages[5]); // Tool result 2
+      expect(result.context[3]).toBe(messages[6]); // AI message 3 with tool use
+      expect(result.context[4]).toBe(messages[7]); // Tool result 3
+    });
+  });
+
   describe('Integration with Run', () => {
     it('should initialize Run with custom token counter and process messages', async () => {
       const provider = Providers.OPENAI;

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

@@ -235,9 +235,9 @@ describe('Token Distribution Edge Case Tests', () => {
     });
 
     // Add two more messages
+    messages.push(new HumanMessage('Message 4'));
     const extendedMessages = [
       ...messages,
-      new HumanMessage('Message 4'),
       new AIMessage('Response 4')
     ];
 
@@ -257,6 +257,7 @@ describe('Token Distribution Edge Case Tests', () => {
     // 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 => {
@@ -282,14 +283,12 @@ describe('Token Distribution Edge Case Tests', () => {
     // 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)) {
-        // This message is not in the context, so its token count should not have been adjusted in the last operation
         const expectedValue = i < messages.length
           ? (secondResult.indexTokenCountMap[i] || indexTokenCountMap[i])
-          : (indexTokenCountMap as Record<string, number | undefined>)[i] ?? indexTokenCountMap[i - 1];
+          : (indexTokenCountMap as Record<string, number | undefined>)[i] ?? 0;
 
-        // For defined values, we can check that they're close to what we expect
         const difference = Math.abs((thirdResult.indexTokenCountMap[i] || 0) - expectedValue);
-        expect(difference).toBeLessThan(20); // Allow for some implementation differences
+        expect(difference).toBe(0);
       }
     }
   });

package/src/types/llm.ts

@@ -35,16 +35,23 @@ export type AzureClientOptions = (Partial<OpenAIChatInput> & Partial<AzureOpenAI
     } & BaseChatModelParams & {
         configuration?: OAIClientOptions;
     });
-
+export type ThinkingConfig = AnthropicInput['thinking'];
 export type ChatOpenAIToolType = BindToolsInput | OpenAIClient.ChatCompletionTool;
 export type CommonToolType = StructuredTool | ChatOpenAIToolType;
-
+export type AnthropicReasoning = {
+  thinking?: ThinkingConfig | boolean;
+  thinkingBudget?: number;
+};
 export type OpenAIClientOptions = ChatOpenAIFields;
 export type OllamaClientOptions = ChatOllamaInput;
 export type AnthropicClientOptions = AnthropicInput;
 export type MistralAIClientOptions = ChatMistralAIInput;
 export type VertexAIClientOptions = ChatVertexAIInput;
 export type BedrockClientOptions = BedrockChatFields;
+export type BedrockAnthropicInput = ChatBedrockConverseInput & {
+  additionalModelRequestFields?: ChatBedrockConverseInput['additionalModelRequestFields'] &
+    AnthropicReasoning;
+};
 export type BedrockConverseClientOptions = ChatBedrockConverseInput;
 export type GoogleClientOptions = GoogleGenerativeAIChatInput;
 export type DeepSeekClientOptions = ChatDeepSeekCallOptions;