illuma-agents 1.0.7 → 1.0.8

package/src/tools/ToolNode.ts

@@ -1,10 +1,17 @@
+import { ToolCall } from '@langchain/core/messages/tool';
+import {
+  ToolMessage,
+  isAIMessage,
+  isBaseMessage,
+} from '@langchain/core/messages';
 import {
   END,
-  MessagesAnnotation,
+  Send,
+  Command,
   isCommand,
   isGraphInterrupt,
+  MessagesAnnotation,
 } from '@langchain/langgraph';
-import { ToolMessage, isBaseMessage } from '@langchain/core/messages';
 import type {
   RunnableConfig,
   RunnableToolLike,
@@ -14,12 +21,20 @@ import type { StructuredToolInterface } from '@langchain/core/tools';
 import type * as t from '@/types';
 import { RunnableCallable } from '@/utils';
 
+/**
+ * Helper to check if a value is a Send object
+ */
+function isSend(value: unknown): value is Send {
+  return value instanceof Send;
+}
+
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 export class ToolNode<T = any> extends RunnableCallable<T, T> {
   tools: t.GenericTool[];
   private toolMap: Map<string, StructuredToolInterface | RunnableToolLike>;
   private loadRuntimeTools?: t.ToolRefGenerator;
   handleToolErrors = true;
+  trace = false;
   toolCallStepIds?: Map<string, string>;
   errorHandler?: t.ToolNodeConstructorParams['errorHandler'];
   private toolUsageCount: Map<string, number>;
@@ -52,60 +67,50 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
     return new Map(this.toolUsageCount); // Return a copy
   }
 
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  protected async run(input: any, config: RunnableConfig): Promise<T> {
-    const message = Array.isArray(input)
-      ? input[input.length - 1]
-      : input.messages[input.messages.length - 1];
-
-    if (message._getType() !== 'ai') {
-      throw new Error('ToolNode only accepts AIMessages as input.');
-    }
-
-    if (this.loadRuntimeTools) {
-      const { tools, toolMap } = this.loadRuntimeTools(
-        (message as AIMessage).tool_calls ?? []
+  /**
+   * Runs a single tool call with error handling
+   */
+  protected async runTool(
+    call: ToolCall,
+    config: RunnableConfig
+  ): Promise<BaseMessage | Command> {
+    const tool = this.toolMap.get(call.name);
+    try {
+      if (tool === undefined) {
+        throw new Error(`Tool "${call.name}" not found.`);
+      }
+      const turn = this.toolUsageCount.get(call.name) ?? 0;
+      this.toolUsageCount.set(call.name, turn + 1);
+      const args = call.args;
+      const stepId = this.toolCallStepIds?.get(call.id!);
+      const output = await tool.invoke(
+        { ...call, args, type: 'tool_call', stepId, turn },
+        config
       );
-      this.tools = tools;
-      this.toolMap = toolMap ?? new Map(tools.map((tool) => [tool.name, tool]));
-    }
-    const outputs = await Promise.all(
-      (message as AIMessage).tool_calls?.map(async (call) => {
-        const tool = this.toolMap.get(call.name);
+      if (
+        (isBaseMessage(output) && output._getType() === 'tool') ||
+        isCommand(output)
+      ) {
+        return output;
+      } else {
+        return new ToolMessage({
+          status: 'success',
+          name: tool.name,
+          content: typeof output === 'string' ? output : JSON.stringify(output),
+          tool_call_id: call.id!,
+        });
+      }
+    } catch (_e: unknown) {
+      const e = _e as Error;
+      if (!this.handleToolErrors) {
+        throw e;
+      }
+      if (isGraphInterrupt(e)) {
+        throw e;
+      }
+      if (this.errorHandler) {
         try {
-          if (tool === undefined) {
-            throw new Error(`Tool "${call.name}" not found.`);
-          }
-          const turn = this.toolUsageCount.get(call.name) ?? 0;
-          this.toolUsageCount.set(call.name, turn + 1);
-          const args = call.args;
-          const stepId = this.toolCallStepIds?.get(call.id!);
-          const output = await tool.invoke(
-            { ...call, args, type: 'tool_call', stepId, turn },
-            config
-          );
-          if (
-            (isBaseMessage(output) && output._getType() === 'tool') ||
-            isCommand(output)
-          ) {
-            return output;
-          } else {
-            return new ToolMessage({
-              name: tool.name,
-              content:
-                typeof output === 'string' ? output : JSON.stringify(output),
-              tool_call_id: call.id!,
-            });
-          }
-        } catch (_e: unknown) {
-          const e = _e as Error;
-          if (!this.handleToolErrors) {
-            throw e;
-          }
-          if (isGraphInterrupt(e)) {
-            throw e;
-          }
-          this.errorHandler?.(
+          await this.errorHandler(
             {
               error: e,
               id: call.id!,
@@ -114,27 +119,153 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
             },
             config.metadata
           );
-          return new ToolMessage({
-            content: `Error: ${e.message}\n Please fix your mistakes.`,
-            name: call.name,
-            tool_call_id: call.id ?? '',
+        } catch (handlerError) {
+          // eslint-disable-next-line no-console
+          console.error('Error in errorHandler:', {
+            toolName: call.name,
+            toolCallId: call.id,
+            toolArgs: call.args,
+            stepId: this.toolCallStepIds?.get(call.id!),
+            turn: this.toolUsageCount.get(call.name),
+            originalError: {
+              message: e.message,
+              stack: e.stack ?? undefined,
+            },
+            handlerError:
+              handlerError instanceof Error
+                ? {
+                  message: handlerError.message,
+                  stack: handlerError.stack ?? undefined,
+                }
+                : {
+                  message: String(handlerError),
+                  stack: undefined,
+                },
           });
         }
-      }) ?? []
-    );
+      }
+      return new ToolMessage({
+        status: 'error',
+        content: `Error: ${e.message}\n Please fix your mistakes.`,
+        name: call.name,
+        tool_call_id: call.id ?? '',
+      });
+    }
+  }
+
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  protected async run(input: any, config: RunnableConfig): Promise<T> {
+    let outputs: (BaseMessage | Command)[];
+
+    if (this.isSendInput(input)) {
+      outputs = [await this.runTool(input.lg_tool_call, config)];
+    } else {
+      let messages: BaseMessage[];
+      if (Array.isArray(input)) {
+        messages = input;
+      } else if (this.isMessagesState(input)) {
+        messages = input.messages;
+      } else {
+        throw new Error(
+          'ToolNode only accepts BaseMessage[] or { messages: BaseMessage[] } as input.'
+        );
+      }
+
+      const toolMessageIds: Set<string> = new Set(
+        messages
+          .filter((msg) => msg._getType() === 'tool')
+          .map((msg) => (msg as ToolMessage).tool_call_id)
+      );
+
+      let aiMessage: AIMessage | undefined;
+      for (let i = messages.length - 1; i >= 0; i--) {
+        const message = messages[i];
+        if (isAIMessage(message)) {
+          aiMessage = message;
+          break;
+        }
+      }
+
+      if (aiMessage == null || !isAIMessage(aiMessage)) {
+        throw new Error('ToolNode only accepts AIMessages as input.');
+      }
+
+      if (this.loadRuntimeTools) {
+        const { tools, toolMap } = this.loadRuntimeTools(
+          aiMessage.tool_calls ?? []
+        );
+        this.tools = tools;
+        this.toolMap =
+          toolMap ?? new Map(tools.map((tool) => [tool.name, tool]));
+      }
+
+      outputs = await Promise.all(
+        aiMessage.tool_calls
+          ?.filter((call) => call.id == null || !toolMessageIds.has(call.id))
+          .map((call) => this.runTool(call, config)) ?? []
+      );
+    }
 
     if (!outputs.some(isCommand)) {
       return (Array.isArray(input) ? outputs : { messages: outputs }) as T;
     }
 
-    const combinedOutputs = outputs.map((output) => {
+    const combinedOutputs: (
+      | { messages: BaseMessage[] }
+      | BaseMessage[]
+      | Command
+    )[] = [];
+    let parentCommand: Command | null = null;
+
+    for (const output of outputs) {
       if (isCommand(output)) {
-        return output;
+        if (
+          output.graph === Command.PARENT &&
+          Array.isArray(output.goto) &&
+          output.goto.every((send): send is Send => isSend(send))
+        ) {
+          if (parentCommand) {
+            (parentCommand.goto as Send[]).push(...(output.goto as Send[]));
+          } else {
+            parentCommand = new Command({
+              graph: Command.PARENT,
+              goto: output.goto,
+            });
+          }
+        } else {
+          combinedOutputs.push(output);
+        }
+      } else {
+        combinedOutputs.push(
+          Array.isArray(input) ? [output] : { messages: [output] }
+        );
       }
-      return Array.isArray(input) ? [output] : { messages: [output] };
-    });
+    }
+
+    if (parentCommand) {
+      combinedOutputs.push(parentCommand);
+    }
+
     return combinedOutputs as T;
   }
+
+  private isSendInput(input: unknown): input is { lg_tool_call: ToolCall } {
+    return (
+      typeof input === 'object' && input != null && 'lg_tool_call' in input
+    );
+  }
+
+  private isMessagesState(
+    input: unknown
+  ): input is { messages: BaseMessage[] } {
+    return (
+      typeof input === 'object' &&
+      input != null &&
+      'messages' in input &&
+      Array.isArray((input as { messages: unknown }).messages) &&
+      (input as { messages: unknown[] }).messages.every(isBaseMessage)
+    );
+  }
 }
 
 function areToolCallsInvoked(

package/src/types/tools.ts

@@ -1,80 +1,80 @@
-// src/types/tools.ts
-import type { StructuredToolInterface } from '@langchain/core/tools';
-import type { RunnableToolLike } from '@langchain/core/runnables';
-import type { ToolCall } from '@langchain/core/messages/tool';
-import type { ToolErrorData } from './stream';
-import { EnvVar } from '@/common';
-
-/** Replacement type for `import type { ToolCall } from '@langchain/core/messages/tool'` in order to have stringified args typed */
-export type CustomToolCall = {
-  name: string;
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  args: string | Record<string, any>;
-  id?: string;
-  type?: 'tool_call';
-  output?: string;
-};
-
-export type GenericTool = StructuredToolInterface | RunnableToolLike;
-
-export type ToolMap = Map<string, GenericTool>;
-export type ToolRefs = {
-  tools: GenericTool[];
-  toolMap?: ToolMap;
-};
-
-export type ToolRefGenerator = (tool_calls: ToolCall[]) => ToolRefs;
-
-export type ToolNodeOptions = {
-  name?: string;
-  tags?: string[];
-  handleToolErrors?: boolean;
-  loadRuntimeTools?: ToolRefGenerator;
-  toolCallStepIds?: Map<string, string>;
-  errorHandler?: (
-    data: ToolErrorData,
-    metadata?: Record<string, unknown>
-  ) => void;
-};
-
-export type ToolNodeConstructorParams = ToolRefs & ToolNodeOptions;
-
-export type ToolEndEvent = {
-  /** The Step Id of the Tool Call */
-  id: string;
-  /** The Completed Tool Call */
-  tool_call: ToolCall;
-  /** The content index of the tool call */
-  index: number;
-};
-
-export type CodeEnvFile = {
-  id: string;
-  name: string;
-  session_id: string;
-};
-
-export type CodeExecutionToolParams =
-  | undefined
-  | {
-      session_id?: string;
-      user_id?: string;
-      apiKey?: string;
-      files?: CodeEnvFile[];
-      [EnvVar.CODE_API_KEY]?: string;
-    };
-
-export type FileRef = {
-  id: string;
-  name: string;
-  path?: string;
-};
-
-export type FileRefs = FileRef[];
-
-export type ExecuteResult = {
-  session_id: string;
-  stdout: string;
-  stderr: string;
-  files?: FileRefs;
-};
+// src/types/tools.ts
+import type { StructuredToolInterface } from '@langchain/core/tools';
+import type { RunnableToolLike } from '@langchain/core/runnables';
+import type { ToolCall } from '@langchain/core/messages/tool';
+import type { ToolErrorData } from './stream';
+import { EnvVar } from '@/common';
+
+/** Replacement type for `import type { ToolCall } from '@langchain/core/messages/tool'` in order to have stringified args typed */
+export type CustomToolCall = {
+  name: string;
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  args: string | Record<string, any>;
+  id?: string;
+  type?: 'tool_call';
+  output?: string;
+};
+
+export type GenericTool = StructuredToolInterface | RunnableToolLike;
+
+export type ToolMap = Map<string, GenericTool>;
+export type ToolRefs = {
+  tools: GenericTool[];
+  toolMap?: ToolMap;
+};
+
+export type ToolRefGenerator = (tool_calls: ToolCall[]) => ToolRefs;
+
+export type ToolNodeOptions = {
+  name?: string;
+  tags?: string[];
+  handleToolErrors?: boolean;
+  loadRuntimeTools?: ToolRefGenerator;
+  toolCallStepIds?: Map<string, string>;
+  errorHandler?: (
+    data: ToolErrorData,
+    metadata?: Record<string, unknown>
+  ) => Promise<void>;
+};
+
+export type ToolNodeConstructorParams = ToolRefs & ToolNodeOptions;
+
+export type ToolEndEvent = {
+  /** The Step Id of the Tool Call */
+  id: string;
+  /** The Completed Tool Call */
+  tool_call: ToolCall;
+  /** The content index of the tool call */
+  index: number;
+};
+
+export type CodeEnvFile = {
+  id: string;
+  name: string;
+  session_id: string;
+};
+
+export type CodeExecutionToolParams =
+  | undefined
+  | {
+      session_id?: string;
+      user_id?: string;
+      apiKey?: string;
+      files?: CodeEnvFile[];
+      [EnvVar.CODE_API_KEY]?: string;
+    };
+
+export type FileRef = {
+  id: string;
+  name: string;
+  path?: string;
+};
+
+export type FileRefs = FileRef[];
+
+export type ExecuteResult = {
+  session_id: string;
+  stdout: string;
+  stderr: string;
+  files?: FileRefs;
+};