@librechat/agents 2.4.322 → 3.0.0-rc1

package/src/run.ts

@@ -1,32 +1,43 @@
 // 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 {
+  createCompletionTitleRunnable,
+  createTitleRunnable,
+} from '@/utils/title';
+import { GraphEvents, Callback, TitleMethod } from '@/common';
+import { MultiAgentGraph } from '@/graphs/MultiAgentGraph';
 import { createTokenCounter } from '@/utils/tokens';
 import { StandardGraph } from '@/graphs/Graph';
 import { HandlerRegistry } from '@/events';
 import { isOpenAILike } from '@/utils/llm';
 
-export class Run<T extends t.BaseGraphState> {
-  graphRunnable?: t.CompiledWorkflow<T, Partial<T>, string>;
-  // private collab!: CollabGraph;
-  // private taskManager!: TaskManager;
-  private handlerRegistry: HandlerRegistry;
+export const defaultOmitOptions = new Set([
+  'stream',
+  'thinking',
+  'streaming',
+  'maxTokens',
+  'clientOptions',
+  'thinkingConfig',
+  'thinkingBudget',
+  'includeThoughts',
+  'maxOutputTokens',
+  'additionalModelRequestFields',
+]);
+
+export class Run<_T extends t.BaseGraphState> {
   id: string;
-  Graph: StandardGraph | undefined;
-  provider: Providers | undefined;
+  private tokenCounter?: t.TokenCounter;
+  private handlerRegistry: HandlerRegistry;
+  private indexTokenCountMap?: Record<string, number>;
+  graphRunnable?: t.CompiledStateWorkflow;
+  Graph: StandardGraph | MultiAgentGraph | undefined;
   returnContent: boolean = false;
 
   private constructor(config: Partial<t.RunConfig>) {
@@ -36,6 +47,8 @@ export class Run<T extends t.BaseGraphState> {
     }
 
     this.id = runId;
+    this.tokenCounter = config.tokenCounter;
+    this.indexTokenCountMap = config.indexTokenCountMap;
 
     const handlerRegistry = new HandlerRegistry();
 
@@ -53,12 +66,18 @@ export class Run<T extends t.BaseGraphState> {
       throw new Error('Graph config not provided');
     }
 
-    if (config.graphConfig.type === 'standard' || !config.graphConfig.type) {
-      this.provider = config.graphConfig.llmConfig.provider;
-      this.graphRunnable = this.createStandardGraph(
-        config.graphConfig
-      ) as unknown as t.CompiledWorkflow<T, Partial<T>, string>;
+    /** Handle different graph types */
+    if (config.graphConfig.type === 'multi-agent') {
+      this.graphRunnable = this.createMultiAgentGraph(config.graphConfig);
+      if (this.Graph) {
+        this.Graph.handlerRegistry = handlerRegistry;
+      }
+    } else {
+      // Default to legacy graph for 'standard' or undefined type
+      this.graphRunnable = this.createLegacyGraph(config.graphConfig);
       if (this.Graph) {
+        this.Graph.compileOptions =
+          config.graphConfig.compileOptions ?? this.Graph.compileOptions;
         this.Graph.handlerRegistry = handlerRegistry;
       }
     }
@@ -66,26 +85,70 @@ export class Run<T extends t.BaseGraphState> {
     this.returnContent = config.returnContent ?? false;
   }
 
-  private createStandardGraph(
-    config: t.StandardGraphConfig
-  ): t.CompiledWorkflow<t.IState, Partial<t.IState>, string> {
-    const { llmConfig, tools = [], ...graphInput } = config;
+  private createLegacyGraph(
+    config: t.LegacyGraphConfig
+  ): t.CompiledStateWorkflow {
+    const {
+      type: _type,
+      llmConfig,
+      signal,
+      tools = [],
+      ...agentInputs
+    } = config;
     const { provider, ...clientOptions } = llmConfig;
 
-    const standardGraph = new StandardGraph({
+    /** TEMP: Create agent configuration for the single agent */
+    const agentConfig: t.AgentInputs = {
+      ...agentInputs,
       tools,
       provider,
       clientOptions,
-      ...graphInput,
+      agentId: 'default',
+    };
+
+    const standardGraph = new StandardGraph({
+      signal,
       runId: this.id,
+      agents: [agentConfig],
+      tokenCounter: this.tokenCounter,
+      indexTokenCountMap: this.indexTokenCountMap,
     });
+    // propagate compile options from graph config
+    standardGraph.compileOptions = (
+      config as t.LegacyGraphConfig
+    ).compileOptions;
     this.Graph = standardGraph;
     return standardGraph.createWorkflow();
   }
 
+  private createMultiAgentGraph(
+    config: t.MultiAgentGraphConfig
+  ): t.CompiledStateWorkflow {
+    const { agents, edges, compileOptions } = config;
+
+    const multiAgentGraph = new MultiAgentGraph({
+      runId: this.id,
+      agents,
+      edges,
+      tokenCounter: this.tokenCounter,
+      indexTokenCountMap: this.indexTokenCountMap,
+    });
+
+    if (compileOptions != null) {
+      multiAgentGraph.compileOptions = compileOptions;
+    }
+
+    this.Graph = multiAgentGraph;
+    return multiAgentGraph.createWorkflow();
+  }
+
   static async create<T extends t.BaseGraphState>(
     config: t.RunConfig
   ): Promise<Run<T>> {
+    // Create tokenCounter if indexTokenCountMap is provided but tokenCounter is not
+    if (config.indexTokenCountMap && !config.tokenCounter) {
+      config.tokenCounter = await createTokenCounter();
+    }
     return new Run<T>(config);
   }
 
@@ -98,12 +161,49 @@ export class Run<T extends t.BaseGraphState> {
     return this.Graph.getRunMessages();
   }
 
+  /**
+   * Creates a custom event callback handler that intercepts custom events
+   * and processes them through our handler registry instead of EventStreamCallbackHandler
+   */
+  private createCustomEventCallback() {
+    return async (
+      eventName: string,
+      data: unknown,
+      runId: string,
+      tags?: string[],
+      metadata?: Record<string, unknown>
+    ): Promise<void> => {
+      if (
+        (data as t.StreamEventData)['emitted'] === true &&
+        eventName === GraphEvents.CHAT_MODEL_STREAM
+      ) {
+        return;
+      }
+      const handler = this.handlerRegistry.getHandler(eventName);
+      if (handler && this.Graph) {
+        await handler.handle(
+          eventName,
+          data as
+            | t.StreamEventData
+            | t.ModelEndData
+            | t.RunStep
+            | t.RunStepDeltaEvent
+            | t.MessageDeltaEvent
+            | t.ReasoningDeltaEvent
+            | { result: t.ToolEndEvent },
+          metadata,
+          this.Graph
+        );
+      }
+    };
+  }
+
   async processStream(
     inputs: t.IState,
     config: Partial<RunnableConfig> & { version: 'v1' | 'v2'; run_id?: string },
     streamOptions?: t.EventStreamOptions
   ): Promise<MessageContentComplex[] | undefined> {
-    if (!this.graphRunnable) {
+    if (this.graphRunnable == null) {
       throw new Error(
         'Run not initialized. Make sure to use Run.create() to instantiate the Run.'
       );
@@ -115,88 +215,45 @@ export class Run<T extends t.BaseGraphState> {
     }
 
     this.Graph.resetValues(streamOptions?.keepContent);
-    const provider = this.Graph.provider;
-    const hasTools = this.Graph.tools ? this.Graph.tools.length > 0 : false;
-    if (streamOptions?.callbacks) {
-      /* TODO: conflicts with callback manager */
-      const callbacks = (config.callbacks as t.ProvidedCallbacks) ?? [];
-      config.callbacks = callbacks.concat(
-        this.getCallbacks(streamOptions.callbacks)
-      );
-    }
+
+    /** Custom event callback to intercept and handle custom events */
+    const customEventCallback = this.createCustomEventCallback();
+
+    const baseCallbacks = (config.callbacks as t.ProvidedCallbacks) ?? [];
+    const streamCallbacks = streamOptions?.callbacks
+      ? this.getCallbacks(streamOptions.callbacks)
+      : [];
+
+    config.callbacks = baseCallbacks.concat(streamCallbacks).concat({
+      [Callback.CUSTOM_EVENT]: customEventCallback,
+    });
 
     if (!this.id) {
       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 as Partial<t.GenericTool>).schema) {
-          return acc;
-        }
-
-        const jsonSchema = zodToJsonSchema(
-          (tool.schema as t.ZodObjectAny).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);
-    }
-    const tokenMap = streamOptions?.indexTokenCountMap ?? {};
-    if (this.Graph.systemMessage && instructionTokens > 0) {
-      this.Graph.indexTokenCountMap = shiftIndexTokenCountMap(
-        tokenMap,
-        instructionTokens
-      );
-    } else if (instructionTokens > 0) {
-      tokenMap[0] = tokenMap[0] + instructionTokens;
-      this.Graph.indexTokenCountMap = tokenMap;
-    } else {
-      this.Graph.indexTokenCountMap = tokenMap;
-    }
-
-    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,
     });
 
-    const stream = this.graphRunnable.streamEvents(inputs, config);
+    const stream = this.graphRunnable.streamEvents(inputs, config, {
+      raiseError: true,
+    });
 
     for await (const event of stream) {
-      const { data, name, metadata, ...info } = event;
+      const { data, metadata, ...info } = event;
 
-      let eventName: t.EventName = info.event;
-      if (
-        hasTools &&
-        manualToolStreamProviders.has(provider) &&
-        eventName === GraphEvents.CHAT_MODEL_STREAM
-      ) {
-        /* Skipping CHAT_MODEL_STREAM event due to double-call edge case */
-        continue;
-      }
+      const eventName: t.EventName = info.event;
 
-      if (eventName && eventName === GraphEvents.ON_CUSTOM_EVENT) {
-        eventName = name;
+      /** Skip custom events as they're handled by our callback */
+      if (eventName === GraphEvents.ON_CUSTOM_EVENT) {
+        continue;
       }
 
       const handler = this.handlerRegistry.getHandler(eventName);
       if (handler) {
-        handler.handle(eventName, data, metadata, this.Graph);
+        await handler.handle(eventName, data, metadata, this.Graph);
       }
     }
 
@@ -205,19 +262,19 @@ export class Run<T extends t.BaseGraphState> {
     }
   }
 
-  private createSystemCallback<K extends keyof ClientCallbacks>(
-    clientCallbacks: ClientCallbacks,
+  private createSystemCallback<K extends keyof t.ClientCallbacks>(
+    clientCallbacks: t.ClientCallbacks,
     key: K
-  ): SystemCallbacks[K] {
+  ): t.SystemCallbacks[K] {
     return ((...args: unknown[]) => {
       const clientCallback = clientCallbacks[key];
       if (clientCallback && this.Graph) {
         (clientCallback as (...args: unknown[]) => void)(this.Graph, ...args);
       }
-    }) as SystemCallbacks[K];
+    }) as t.SystemCallbacks[K];
   }
 
-  getCallbacks(clientCallbacks: ClientCallbacks): SystemCallbacks {
+  getCallbacks(clientCallbacks: t.ClientCallbacks): t.SystemCallbacks {
     return {
       [Callback.TOOL_ERROR]: this.createSystemCallback(
         clientCallbacks,
@@ -235,15 +292,18 @@ export class Run<T extends t.BaseGraphState> {
   }
 
   async generateTitle({
+    provider,
     inputText,
     contentParts,
     titlePrompt,
     clientOptions,
     chainOptions,
     skipLanguage,
-  }: t.RunTitleOptions): Promise<{ language: string; title: string }> {
+    titleMethod = TitleMethod.COMPLETION,
+    titlePromptTemplate,
+  }: t.RunTitleOptions): Promise<{ language?: string; title?: string }> {
     const convoTemplate = PromptTemplate.fromTemplate(
-      'User: {input}\nAI: {output}'
+      titlePromptTemplate ?? 'User: {input}\nAI: {output}'
     );
     const response = contentParts
       .map((part) => {
@@ -255,22 +315,14 @@ export class Run<T extends t.BaseGraphState> {
       await convoTemplate.invoke({ input: inputText, output: response })
     ).value;
     const model = this.Graph?.getNewModel({
+      provider,
       clientOptions,
-      omitOriginalOptions: new Set([
-        'clientOptions',
-        'streaming',
-        'stream',
-        'thinking',
-        'maxTokens',
-        'maxOutputTokens',
-        'additionalModelRequestFields',
-      ]),
     });
     if (!model) {
       return { language: '', title: '' };
     }
     if (
-      isOpenAILike(this.provider) &&
+      isOpenAILike(provider) &&
       (model instanceof ChatOpenAI || model instanceof AzureChatOpenAI)
     ) {
       model.temperature = (clientOptions as t.OpenAIClientOptions | undefined)
@@ -286,10 +338,27 @@ export class Run<T extends t.BaseGraphState> {
       model.n = (clientOptions as t.OpenAIClientOptions | undefined)
         ?.n as number;
     }
-    const chain = await createTitleRunnable(model, titlePrompt);
-    return (await chain.invoke(
-      { convo, inputText, skipLanguage },
-      chainOptions
-    )) as { language: string; title: string };
+    const chain =
+      titleMethod === TitleMethod.COMPLETION
+        ? await createCompletionTitleRunnable(model, titlePrompt)
+        : await createTitleRunnable(model, titlePrompt);
+    const invokeConfig = Object.assign({}, chainOptions, {
+      run_id: this.id,
+      runId: this.id,
+    });
+    try {
+      return await chain.invoke(
+        { convo, inputText, skipLanguage },
+        invokeConfig
+      );
+    } catch (_e) {
+      // Fallback: strip callbacks to avoid EventStream tracer errors in certain environments
+      const { callbacks: _cb, ...rest } = invokeConfig;
+      const safeConfig = Object.assign({}, rest, { callbacks: [] });
+      return await chain.invoke(
+        { convo, inputText, skipLanguage },
+        safeConfig as Partial<RunnableConfig>
+      );
+    }
   }
 }

package/src/scripts/ant_web_search.ts

@@ -0,0 +1,158 @@
+/* eslint-disable no-console */
+// src/scripts/cli.ts
+import { config } from 'dotenv';
+config();
+import { HumanMessage, BaseMessage } from '@langchain/core/messages';
+import type * as t from '@/types';
+import { ChatModelStreamHandler, createContentAggregator } from '@/stream';
+import { ToolEndHandler, ModelEndHandler } from '@/events';
+
+import { getArgs } from '@/scripts/args';
+import { Run } from '@/run';
+import { GraphEvents, Callback, Providers } from '@/common';
+import { getLLMConfig } from '@/utils/llmConfig';
+
+const conversationHistory: BaseMessage[] = [];
+let _contentParts: (t.MessageContentComplex | undefined)[] = [];
+async function testStandardStreaming(): Promise<void> {
+  const { userName, location, currentDate } = await getArgs();
+  const { contentParts, aggregateContent } = createContentAggregator();
+  _contentParts = contentParts;
+  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(
+    Providers.ANTHROPIC
+  ) as t.AnthropicClientOptions & t.SharedLLMConfig;
+  llmConfig.model = 'claude-3-5-sonnet-latest';
+
+  const run = await Run.create<t.IState>({
+    runId: 'test-run-id',
+    graphConfig: {
+      type: 'standard',
+      llmConfig,
+      tools: [
+        {
+          type: 'web_search_20250305',
+          name: 'web_search',
+          max_uses: 5,
+        },
+      ],
+      instructions: 'You are a friendly AI assistant.',
+      // additional_instructions: `Always address the user by their name. The user's name is ${userName} and they are located in ${location}.`,
+    },
+    returnContent: true,
+    customHandlers,
+  });
+
+  const config = {
+    configurable: {
+      provider: Providers.ANTHROPIC,
+      thread_id: 'conversation-num-1',
+    },
+    streamMode: 'values',
+    version: 'v2' as const,
+  };
+
+  console.log('Test 1: Search query (search tool test)');
+
+  // const userMessage = `
+  // Make a search for the weather in ${location} today, which is ${currentDate}.
+  // Before making the search, please let me know what you're about to do, then immediately start searching without hesitation.
+  // Make sure to always refer to me by name, which is ${userName}.
+  // After giving me a thorough summary, tell me a joke about the weather forecast we went over.
+  // `;
+  // const userMessage = 'Are massage guns good?';
+  // const userMessage = 'What is functional programming?';
+  const userMessage = "Get me today's trending news.";
+  // const userMessage = "search recent italy earthquake volcano activity";
+  // const userMessage =
+  //   "use 'Trump' as the exact search query and tell me what you find.";
+
+  conversationHistory.push(new HumanMessage(userMessage));
+
+  const inputs = {
+    messages: conversationHistory,
+  };
+  const finalContentParts = await run.processStream(inputs, config);
+  const finalMessages = run.getRunMessages();
+  if (finalMessages) {
+    conversationHistory.push(...finalMessages);
+    console.dir(conversationHistory, { depth: null });
+  }
+  // console.dir(finalContentParts, { depth: null });
+  console.log('\n\n====================\n\n');
+  // console.dir(contentParts, { depth: null });
+}
+
+process.on('unhandledRejection', (reason, promise) => {
+  console.error('Unhandled Rejection at:', promise, 'reason:', reason);
+  console.log('Content Parts:');
+  console.dir(_contentParts, { depth: null });
+  process.exit(1);
+});
+
+testStandardStreaming().catch((err) => {
+  console.error(err);
+  console.log('Conversation history:');
+  console.dir(conversationHistory, { depth: null });
+  console.log('Content Parts:');
+  console.dir(_contentParts, { depth: null });
+  process.exit(1);
+});

package/src/scripts/args.ts

@@ -4,41 +4,45 @@ import { hideBin } from 'yargs/helpers';
 import { llmConfigs } from '@/utils/llmConfig';
 import { Providers } from '@/common';
 
-export async function getArgs(): Promise<{ userName: string; location: string; provider: string; currentDate: string; }> {
+export async function getArgs(): Promise<{
+  userName: string;
+  location: string;
+  provider: Providers;
+  currentDate: string;
+}> {
   const argv = yargs(hideBin(process.argv))
     .option('name', {
       alias: 'n',
       type: 'string',
       description: 'User name',
-      default: 'Jo'
+      default: 'Jo',
     })
     .option('location', {
       alias: 'l',
       type: 'string',
       description: 'User location',
-      default: 'New York'
+      default: 'New York',
     })
     .option('provider', {
       alias: 'p',
       type: 'string',
       description: 'LLM provider',
       choices: Object.keys(llmConfigs),
-      default: Providers.OPENAI
+      default: Providers.OPENAI,
     })
     .help()
-    .alias('help', 'h')
-    .argv;
+    .alias('help', 'h').argv;
 
   const args = await argv;
   const userName = args.name as string;
   const location = args.location as string;
-  const provider = args.provider as string;
+  const provider = args.provider as Providers;
   const currentDate = new Date().toLocaleString();
 
   return {
     userName,
     location,
     provider,
-    currentDate
+    currentDate,
   };
 }