@librechat/agents 2.2.8 → 2.3.0

package/src/messages/prune.ts

@@ -1,3 +1,4 @@
+import { AIMessage } from '@langchain/core/messages';
 import type { BaseMessage, UsageMetadata } from '@langchain/core/messages';
 import type { TokenCounter } from '@/types/run';
 export type PruneMessagesFactoryParams = {
@@ -5,6 +6,7 @@ export type PruneMessagesFactoryParams = {
   startIndex: number;
   tokenCounter: TokenCounter;
   indexTokenCountMap: Record<string, number>;
+  thinkingEnabled?: boolean;
 };
 export type PruneMessagesParams = {
   messages: BaseMessage[];
@@ -45,11 +47,15 @@ function getMessagesWithinTokenLimit({
   maxContextTokens,
   indexTokenCountMap,
   startOnMessageType,
+  thinkingEnabled,
+  tokenCounter,
 }: {
   messages: BaseMessage[];
   maxContextTokens: number;
   indexTokenCountMap: Record<string, number>;
   startOnMessageType?: string;
+  thinkingEnabled?: boolean;
+  tokenCounter?: TokenCounter;
 }): {
   context: BaseMessage[];
   remainingContextTokens: number;
@@ -87,19 +93,204 @@ function getMessagesWithinTokenLimit({
       }
     }
     
-    if (startOnMessageType && context.length > 0) {
-      const requiredTypeIndex = context.findIndex(msg => msg.getType() === startOnMessageType);
-      
-      if (requiredTypeIndex > 0) {
-        context = context.slice(requiredTypeIndex);
-      }
+  // Handle startOnMessageType requirement
+  if (startOnMessageType && context.length > 0) {
+    const requiredTypeIndex = context.findIndex(msg => msg.getType() === startOnMessageType);
+    
+    if (requiredTypeIndex > 0) {
+      context = context.slice(requiredTypeIndex);
     }
   }
-
+  
+  // Add system message if it exists
   if (instructions && _messages.length > 0) {
     context.push(_messages[0] as BaseMessage);
     messages.shift();
   }
+  
+  // Handle thinking mode requirement for Anthropic
+  if (thinkingEnabled && context.length > 0 && tokenCounter) {
+    // Check if the latest message is an assistant message
+    const latestMessageIsAssistant = _messages.length > 0 && _messages[_messages.length - 1].getType() === 'ai';
+    
+    // Process only if we have an assistant message in the context
+    const firstAssistantIndex = context.findIndex(msg => msg.getType() === 'ai');
+    
+    if (firstAssistantIndex >= 0) {
+      const firstAssistantMsg = context[firstAssistantIndex];
+      
+      // Check if the first assistant message already has a thinking block
+      const hasThinkingBlock = Array.isArray(firstAssistantMsg.content) && 
+        firstAssistantMsg.content.some(item => 
+          item && typeof item === 'object' && item.type === 'thinking');
+      
+      // Only proceed if we need to add thinking blocks
+      if (!hasThinkingBlock) {
+        // Collect thinking blocks from pruned assistant messages, starting from the most recent
+        const thinkingBlocks: any[] = [];
+        
+        // Look through pruned messages for thinking blocks, starting from the end (most recent)
+        for (let i = messages.length - 1; i >= 0; i--) {
+          const msg = messages[i];
+          if (msg.getType() === 'ai' && Array.isArray(msg.content)) {
+            for (const item of msg.content) {
+              if (item && typeof item === 'object' && item.type === 'thinking') {
+                thinkingBlocks.push(item);
+                // We only need one thinking block
+                break;
+              }
+            }
+            if (thinkingBlocks.length > 0) break; // Stop after finding one thinking block
+          }
+        }
+        
+        // If we found thinking blocks, add them to the first assistant message
+        if (thinkingBlocks.length > 0) {
+          // Calculate token count of original message
+          const originalTokenCount = tokenCounter(firstAssistantMsg);
+          
+          // Create a new content array with thinking blocks at the beginning
+          let newContent: any[];
+          
+          if (Array.isArray(firstAssistantMsg.content)) {
+            // Keep the original content (excluding any existing thinking blocks)
+            const originalContent = firstAssistantMsg.content.filter(item => 
+              !(item && typeof item === 'object' && item.type === 'thinking'));
+            
+            newContent = [...thinkingBlocks, ...originalContent];
+          } else if (typeof firstAssistantMsg.content === 'string') {
+            newContent = [
+              ...thinkingBlocks,
+              { type: 'text', text: firstAssistantMsg.content }
+            ];
+          } else {
+            newContent = thinkingBlocks;
+          }
+          
+          // Create a new message with the updated content
+          const newMessage = new AIMessage({
+            content: newContent,
+            additional_kwargs: firstAssistantMsg.additional_kwargs,
+            response_metadata: firstAssistantMsg.response_metadata,
+          });
+          
+          // Calculate token count of new message
+          const newTokenCount = tokenCounter(newMessage);
+          
+          // Adjust current token count
+          currentTokenCount += (newTokenCount - originalTokenCount);
+          
+          // Replace the first assistant message
+          context[firstAssistantIndex] = newMessage;
+          
+          // If we've exceeded the token limit, we need to prune more messages
+          if (currentTokenCount > remainingContextTokens) {
+            // Build a map of tool call IDs to track AI <--> tool message correspondences
+            const toolCallIdMap = new Map<string, number>();
+            
+            // Identify tool call IDs in the context
+            for (let i = 0; i < context.length; i++) {
+              const msg = context[i];
+              
+              // Check for tool calls in AI messages
+              if (msg.getType() === 'ai' && Array.isArray(msg.content)) {
+                for (const item of msg.content) {
+                  if (item && typeof item === 'object' && item.type === 'tool_use' && item.id) {
+                    toolCallIdMap.set(item.id, i);
+                  }
+                }
+              }
+              
+              // Check for tool messages
+              if (msg.getType() === 'tool' && 'tool_call_id' in msg && typeof msg.tool_call_id === 'string') {
+                toolCallIdMap.set(msg.tool_call_id, i);
+              }
+            }
+            
+            // Track which messages to remove
+            const indicesToRemove = new Set<number>();
+            
+            // Start removing messages from the end, but preserve AI <--> tool message correspondences
+            let i = context.length - 1;
+            while (i > firstAssistantIndex && currentTokenCount > remainingContextTokens) {
+              const msgToRemove = context[i];
+              
+              // Check if this is a tool message or has tool calls
+              let canRemove = true;
+              
+              if (msgToRemove.getType() === 'tool' && 'tool_call_id' in msgToRemove && typeof msgToRemove.tool_call_id === 'string') {
+                // If this is a tool message, check if we need to keep its corresponding AI message
+                const aiIndex = toolCallIdMap.get(msgToRemove.tool_call_id);
+                if (aiIndex !== undefined && aiIndex !== i && !indicesToRemove.has(aiIndex)) {
+                  // We need to remove both the tool message and its corresponding AI message
+                  indicesToRemove.add(i);
+                  indicesToRemove.add(aiIndex);
+                  currentTokenCount -= (tokenCounter(msgToRemove) + tokenCounter(context[aiIndex]));
+                  canRemove = false;
+                }
+              } else if (msgToRemove.getType() === 'ai' && Array.isArray(msgToRemove.content)) {
+                // If this is an AI message with tool calls, check if we need to keep its corresponding tool messages
+                for (const item of msgToRemove.content) {
+                  if (item && typeof item === 'object' && item.type === 'tool_use' && item.id) {
+                    const toolIndex = toolCallIdMap.get(item.id as string);
+                    if (toolIndex !== undefined && toolIndex !== i && !indicesToRemove.has(toolIndex)) {
+                      // We need to remove both the AI message and its corresponding tool message
+                      indicesToRemove.add(i);
+                      indicesToRemove.add(toolIndex);
+                      currentTokenCount -= (tokenCounter(msgToRemove) + tokenCounter(context[toolIndex]));
+                      canRemove = false;
+                      break;
+                    }
+                  }
+                }
+              }
+              
+              // If we can remove this message individually
+              if (canRemove && !indicesToRemove.has(i)) {
+                indicesToRemove.add(i);
+                currentTokenCount -= tokenCounter(msgToRemove);
+              }
+              
+              i--;
+            }
+            
+            // Remove messages in reverse order to avoid index shifting
+            const sortedIndices = Array.from(indicesToRemove).sort((a, b) => b - a);
+            for (const index of sortedIndices) {
+              context.splice(index, 1);
+            }
+            
+            // Update remainingContextTokens to reflect the new token count
+            remainingContextTokens = maxContextTokens - currentTokenCount;
+          }
+        }
+      }
+    }
+    
+    // If the latest message is an assistant message, ensure an assistant message appears early in the context
+    // but maintain system message precedence
+    if (latestMessageIsAssistant && context.length > 0) {
+      // Find the first assistant message in the context
+      const assistantIndex = context.findIndex(msg => msg.getType() === 'ai');
+      if (assistantIndex > 0) {
+        // Check if there's a system message at the beginning
+        const hasSystemFirst = context[0].getType() === 'system';
+        
+        // Move the assistant message to the appropriate position
+        const assistantMsg = context[assistantIndex];
+        context.splice(assistantIndex, 1);
+        
+        if (hasSystemFirst) {
+          // Insert after the system message
+          context.splice(1, 0, assistantMsg);
+        } else {
+          // Insert at the beginning if no system message
+          context.unshift(assistantMsg);
+        }
+      }
+    }
+  }
+  }
 
   const prunedMemory = messages;
   summaryIndex = prunedMemory.length - 1;
@@ -121,6 +312,7 @@ 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>;
@@ -167,11 +359,14 @@ export function createPruneMessages(factoryParams: PruneMessagesFactoryParams) {
       return { context: params.messages, indexTokenCountMap };
     }
 
+    // Pass the tokenCounter to getMessagesWithinTokenLimit for token recalculation
     const { context } = getMessagesWithinTokenLimit({
       maxContextTokens: factoryParams.maxTokens,
       messages: params.messages,
       indexTokenCountMap,
       startOnMessageType: params.startOnMessageType,
+      thinkingEnabled: factoryParams.thinkingEnabled,
+      tokenCounter: factoryParams.tokenCounter,
     });
 
     return { context, indexTokenCountMap };

package/src/scripts/code_exec_files.ts

@@ -0,0 +1,162 @@
+// src/scripts/cli.ts
+import { config } from 'dotenv';
+config();
+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, createMetadataAggregator } from '@/events';
+import { getLLMConfig } from '@/utils/llmConfig';
+import { getArgs } from '@/scripts/args';
+import { GraphEvents } from '@/common';
+import { Run } from '@/run';
+import { createCodeExecutionTool } from '@/tools/CodeExecutor';
+
+const conversationHistory: BaseMessage[] = [];
+
+async function testCodeExecution(): Promise<void> {
+  const { userName, location, provider, currentDate } = await getArgs();
+  const { contentParts, aggregateContent } = createContentAggregator();
+  const customHandlers = {
+    [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 => {
+        console.log('====== ON_RUN_STEP_COMPLETED ======');
+        console.dir(data, { depth: null });
+        aggregateContent({ event, data: data as unknown as { result: t.ToolEndEvent } });
+      }
+    },
+    [GraphEvents.ON_RUN_STEP]: {
+      handle: (event: GraphEvents.ON_RUN_STEP, data: t.StreamEventData): void => {
+        console.log('====== ON_RUN_STEP ======');
+        console.dir(data, { depth: null });
+        aggregateContent({ event, data: data as t.RunStep });
+      }
+    },
+    [GraphEvents.ON_RUN_STEP_DELTA]: {
+      handle: (event: GraphEvents.ON_RUN_STEP_DELTA, data: t.StreamEventData): void => {
+        console.log('====== ON_RUN_STEP_DELTA ======');
+        console.dir(data, { depth: null });
+        aggregateContent({ event, data: data as t.RunStepDeltaEvent });
+      }
+    },
+    [GraphEvents.ON_MESSAGE_DELTA]: {
+      handle: (event: GraphEvents.ON_MESSAGE_DELTA, data: t.StreamEventData): void => {
+        console.log('====== ON_MESSAGE_DELTA ======');
+        console.dir(data, { depth: null });
+        aggregateContent({ event, data: data as t.MessageDeltaEvent });
+      }
+    },
+    [GraphEvents.TOOL_START]: {
+      handle: (_event: string, data: t.StreamEventData, metadata?: Record<string, unknown>): void => {
+        console.log('====== TOOL_START ======');
+        console.dir(data, { depth: null });
+      }
+    },
+  };
+
+  const llmConfig = getLLMConfig(provider);
+
+  const run = await Run.create<t.IState>({
+    runId: 'message-num-1',
+    graphConfig: {
+      type: 'standard',
+      llmConfig,
+      tools: [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}. The current date is ${currentDate}.`,
+    },
+    returnContent: true,
+    customHandlers,
+  });
+
+  const config: Partial<RunnableConfig> & { version: 'v1' | 'v2'; run_id?: string; streamMode: string } = {
+    configurable: {
+      provider,
+      thread_id: 'conversation-num-1',
+    },
+    streamMode: 'values',
+    version: 'v2' as const,
+    // recursionLimit: 3,
+  };
+
+  console.log('Test 1: Create Project Plan');
+
+  const userMessage1 = `
+  Hi ${userName} here. We are testing your file capabilities.
+  
+  1. Create a text file named "project_plan.txt" that contains: "This is a project plan for a new software development project."
+  
+  Please generate this file so I can review it.
+  `;
+
+  conversationHistory.push(new HumanMessage(userMessage1));
+
+  let inputs = {
+    messages: conversationHistory,
+  };
+  const finalContentParts1 = await run.processStream(inputs, config);
+  const finalMessages1 = run.getRunMessages();
+  if (finalMessages1) {
+    conversationHistory.push(...finalMessages1);
+  }
+  console.log('\n\n====================\n\n');
+  console.dir(contentParts, { depth: null });
+
+  console.log('Test 2: Edit Project Plan');
+
+  const userMessage2 = `
+  Thanks for creating the project plan. Now I'd like you to edit the same plan to:
+  
+  1. Add a new section called "Technology Stack" that contains: "The technology stack for this project includes the following technologies" and nothing more.
+  
+`;
+
+  // Make sure to pass the file ID of the previous file you created and explicitly duplicate or rename the file in your code so we can then access it. Also print the contents of the new file to ensure we did what we wanted.`;
+
+  conversationHistory.push(new HumanMessage(userMessage2));
+
+  inputs = {
+    messages: conversationHistory,
+  };
+  const finalContentParts2 = await run.processStream(inputs, config);
+  const finalMessages2 = run.getRunMessages();
+  if (finalMessages2) {
+    conversationHistory.push(...finalMessages2);
+  }
+  console.log('\n\n====================\n\n');
+  console.dir(contentParts, { depth: null });
+
+  const { handleLLMEnd, collected } = createMetadataAggregator();
+  const titleResult = await run.generateTitle({
+    inputText: userMessage2,
+    contentParts,
+    chainOptions: {
+      callbacks: [{
+        handleLLMEnd,
+      }],
+    },
+  });
+  console.log('Generated Title:', titleResult);
+  console.log('Collected metadata:', collected);
+}
+
+process.on('unhandledRejection', (reason, promise) => {
+  console.error('Unhandled Rejection at:', promise, 'reason:', reason);
+  console.log('Conversation history:');
+  console.dir(conversationHistory, { depth: null });
+  process.exit(1);
+});
+
+process.on('uncaughtException', (err) => {
+  console.error('Uncaught Exception:', err);
+});
+
+testCodeExecution().catch((err) => {
+  console.error(err);
+  console.log('Conversation history:');
+  console.dir(conversationHistory, { depth: null });
+  process.exit(1);
+});