feat: add llm.exec (#2078)

.changeset/slimy-ways-travel.md

@@ -0,0 +1,7 @@
+---
+"@llamaindex/core": patch
+"@llamaindex/tools": patch
+"@llamaindex/workflow": patch
+---
+
+feat: add llm.exec

apps/next/src/app/api/chat/route.ts

@@ -1,4 +1,4 @@
-import { MockLLM } from "@llamaindex/core/utils";
+import { MockLLM } from "@llamaindex/core/llms/mock";
 import { LlamaIndexAdapter, type Message } from "ai";
 import { Settings, SimpleChatEngine, type ChatMessage } from "llamaindex";
 import { NextResponse, type NextRequest } from "next/server";

examples/agents/agent/single-agent.ts

@@ -3,7 +3,7 @@
  */
 import { openai } from "@llamaindex/openai";
 import { agent } from "@llamaindex/workflow";
-import { getWeatherTool } from "../../deprecated/agents/utils/tools";
+import { getWeatherTool } from "../tools/tools";
 
 async function main() {
   const weatherAgent = agent({

examples/agents/agent/with-ollama.ts

@@ -1,6 +1,6 @@
 import { ollama } from "@llamaindex/ollama";
 import { agent } from "@llamaindex/workflow";
-import { getWeatherTool } from "../../deprecated/agents/utils/tools";
+import { getWeatherTool } from "../tools/tools";
 
 async function main() {
   const myAgent = agent({

examples/agents/tools/chat-stream.ts

@@ -1,7 +1,7 @@
-import { OpenAI } from "@llamaindex/openai";
+import { openai } from "@llamaindex/openai";
 
 async function main() {
-  const llm = new OpenAI({ model: "gpt-4-turbo" });
+  const llm = openai({ model: "gpt-4.1-mini" });
   const args: Parameters<typeof llm.chat>[0] = {
     additionalChatOptions: {
       tool_choice: "auto",

examples/agents/tools/exec-stream.ts

@@ -0,0 +1,46 @@
+import { openai } from "@llamaindex/openai";
+import { tool } from "llamaindex";
+import z from "zod";
+
+import { ChatMessage } from "llamaindex";
+
+async function main() {
+  const llm = openai({ model: "gpt-4.1-mini" });
+  const messages = [
+    {
+      content: `What's the weather like in San Francisco?`,
+      role: "user",
+    } as ChatMessage,
+  ];
+
+  let exit = false;
+  do {
+    const { stream, newMessages, toolCalls } = await llm.exec({
+      messages,
+      tools: [
+        tool({
+          name: "get_weather",
+          description: "Get the current weather for a location",
+          parameters: z.object({
+            address: z.string().describe("The address"),
+          }),
+          execute: ({ address }) => {
+            return `It's sunny in ${address}!`;
+          },
+        }),
+      ],
+      stream: true,
+    });
+    for await (const chunk of stream) {
+      process.stdout.write(chunk.delta);
+    }
+    messages.push(...newMessages());
+    // exit condition to stop the agent loop
+    // here we can also check for specific tool calls or limit the number of llm.exec calls
+    exit = toolCalls.length === 0;
+  } while (!exit);
+}
+
+(async function () {
+  await main();
+})();

examples/agents/tools/exec.ts

@@ -0,0 +1,43 @@
+import { openai } from "@llamaindex/openai";
+import { ChatMessage, tool } from "llamaindex";
+import z from "zod";
+
+async function main() {
+  const llm = openai({ model: "gpt-4.1-mini" });
+  const messages = [
+    {
+      content: `What's the weather like in San Francisco?`,
+      role: "user",
+    } as ChatMessage,
+  ];
+
+  let exit = false;
+  do {
+    const { newMessages, toolCalls } = await llm.exec({
+      messages,
+      tools: [
+        tool({
+          name: "get_weather",
+          description: "Get the current weather for a location",
+          parameters: z.object({
+            address: z.string().describe("The address"),
+          }),
+          execute: ({ address }) => {
+            return `It's sunny in ${address}!`;
+          },
+        }),
+      ],
+    });
+    console.log(newMessages);
+    messages.push(...newMessages);
+    // exit condition to stop the agent loop
+    // here we can also check for specific tool calls or limit the number of llm.exec calls
+    exit = toolCalls.length === 0;
+  } while (!exit);
+}
+
+(async function () {
+  console.log("Starting...");
+  await main();
+  console.log("Done");
+})();

examples/deprecated/agents/openai-task.ts

@@ -4,7 +4,7 @@ import {
   getCurrentIDTool,
   getUserInfoTool,
   getWeatherTool,
-} from "./utils/tools";
+} from "../../agents/tools/tools";
 
 async function main() {
   // Create an OpenAIAgent with the function tools

examples/deprecated/agents/react-task.ts

@@ -3,7 +3,7 @@ import {
   getCurrentIDTool,
   getUserInfoTool,
   getWeatherTool,
-} from "./utils/tools";
+} from "../../agents/tools/tools";
 
 async function main() {
   // Create an OpenAIAgent with the function tools

packages/core/package.json

@@ -59,6 +59,17 @@
       },
       "default": "./llms/dist/index.js"
     },
+    "./llms/mock": {
+      "require": {
+        "types": "./llms/mock/dist/index.d.cts",
+        "default": "./llms/mock/dist/index.cjs"
+      },
+      "import": {
+        "types": "./llms/mock/dist/index.d.ts",
+        "default": "./llms/mock/dist/index.js"
+      },
+      "default": "./llms/mock/dist/index.js"
+    },
     "./decorator": {
       "require": {
         "types": "./decorator/dist/index.d.cts",

packages/core/src/llms/base.ts

@@ -1,15 +1,20 @@
 import { extractText } from "../utils/llms";
 import { streamConverter } from "../utils/stream";
+import { callTool, getToolCallsFromResponse } from "./tool-call";
 import type {
+  ChatMessage,
   ChatResponse,
   ChatResponseChunk,
   CompletionResponse,
+  ExecResponse,
+  ExecStreamResponse,
   LLM,
   LLMChatParamsNonStreaming,
   LLMChatParamsStreaming,
   LLMCompletionParamsNonStreaming,
   LLMCompletionParamsStreaming,
   LLMMetadata,
+  PartialToolCall,
   ToolCallLLMMessageOptions,
 } from "./type";
 
@@ -60,13 +65,180 @@ export abstract class BaseLLM<
       AdditionalChatOptions,
       AdditionalMessageOptions
     >,
-  ): Promise<AsyncIterable<ChatResponseChunk>>;
+  ): Promise<AsyncIterable<ChatResponseChunk<AdditionalMessageOptions>>>;
   abstract chat(
     params: LLMChatParamsNonStreaming<
       AdditionalChatOptions,
       AdditionalMessageOptions
     >,
   ): Promise<ChatResponse<AdditionalMessageOptions>>;
+
+  exec(
+    params: LLMChatParamsStreaming<
+      AdditionalChatOptions,
+      AdditionalMessageOptions
+    >,
+  ): Promise<ExecStreamResponse<AdditionalMessageOptions>>;
+  exec(
+    params: LLMChatParamsNonStreaming<
+      AdditionalChatOptions,
+      AdditionalMessageOptions
+    >,
+  ): Promise<ExecResponse<AdditionalMessageOptions>>;
+  async exec(
+    params:
+      | LLMChatParamsStreaming<AdditionalChatOptions, AdditionalMessageOptions>
+      | LLMChatParamsNonStreaming<
+          AdditionalChatOptions,
+          AdditionalMessageOptions
+        >,
+  ): Promise<
+    | ExecResponse<AdditionalMessageOptions>
+    | ExecStreamResponse<AdditionalMessageOptions>
+  > {
+    if (params.stream) {
+      return this.streamExec(params);
+    }
+    const newMessages: ChatMessage<AdditionalMessageOptions>[] = [];
+    const response = await this.chat(params);
+    newMessages.push(response.message);
+    const toolCalls = getToolCallsFromResponse(response);
+    if (params.tools && toolCalls.length > 0) {
+      for (const toolCall of toolCalls) {
+        const toolResultMessage = await callTool<AdditionalMessageOptions>(
+          params.tools,
+          toolCall,
+        );
+        if (toolResultMessage) {
+          newMessages.push(toolResultMessage);
+        }
+      }
+    }
+    return {
+      newMessages,
+      toolCalls,
+    };
+  }
+
+  async streamExec(
+    params: LLMChatParamsStreaming<
+      AdditionalChatOptions,
+      AdditionalMessageOptions
+    >,
+  ): Promise<ExecStreamResponse<AdditionalMessageOptions>> {
+    const responseStream = await this.chat(params);
+    const iterator = responseStream[Symbol.asyncIterator]();
+    const first = await iterator.next();
+
+    // Set firstChunk to null if empty
+    const firstChunk = !first.done ? first.value : null;
+
+    const hasToolCallsInFirst =
+      firstChunk?.options && "toolCall" in firstChunk.options;
+
+    if (!hasToolCallsInFirst) {
+      let content = firstChunk?.delta ?? "";
+      let finished = false;
+      return {
+        stream: (async function* () {
+          if (firstChunk) {
+            yield firstChunk;
+          }
+          for await (const chunk of {
+            [Symbol.asyncIterator]: () => iterator,
+          }) {
+            content += chunk.delta;
+            yield chunk;
+          }
+          finished = true;
+        })(),
+        toolCalls: [],
+        newMessages() {
+          if (!finished) {
+            throw new Error(
+              "New messages are not ready yet. Call newMessages() after the stream is done.",
+            );
+          }
+          return content
+            ? [
+                {
+                  role: "assistant",
+                  content,
+                } as ChatMessage<AdditionalMessageOptions>,
+              ]
+            : [];
+        },
+      };
+    }
+    // Helper function to process a chunk
+    function processChunk(
+      chunk: ChatResponseChunk,
+      toolCallMap: Map<string, PartialToolCall>,
+    ): ChatResponseChunk | null {
+      if (chunk.options && "toolCall" in chunk.options) {
+        // update tool call map
+        for (const toolCall of chunk.options.toolCall as PartialToolCall[]) {
+          if (toolCall.id) {
+            toolCallMap.set(toolCall.id, toolCall);
+          }
+        }
+        // return the current full response with the tool calls
+        const toolCalls = Array.from(toolCallMap.values());
+        return {
+          ...chunk,
+          options: {
+            ...chunk.options,
+            toolCall: toolCalls,
+          },
+        };
+      }
+      return null;
+    }
+    // Collect for tool call
+    let fullResponse: ChatResponseChunk | null = null;
+    const toolCallMap = new Map<string, PartialToolCall>();
+    // Process first chunk
+    fullResponse = processChunk(firstChunk, toolCallMap);
+    // Process remaining chunks
+    while (true) {
+      const next = await iterator.next();
+      if (next.done) break;
+      const chunk = next.value;
+      const potentialFull = processChunk(chunk, toolCallMap);
+      if (potentialFull) {
+        fullResponse = potentialFull;
+      }
+    }
+    if (params.tools && fullResponse) {
+      const toolCalls = getToolCallsFromResponse(fullResponse);
+      const messages: ChatMessage<AdditionalMessageOptions>[] = [];
+      messages.push({
+        role: "assistant",
+        content: "",
+        options: {
+          toolCall: toolCalls,
+        } as AdditionalMessageOptions,
+      });
+      for (const toolCall of toolCalls) {
+        const toolResultMessage = await callTool<AdditionalMessageOptions>(
+          params.tools,
+          toolCall,
+        );
+        if (toolResultMessage) {
+          messages.push(toolResultMessage);
+        }
+      }
+      return {
+        stream: (async function* () {})(),
+        newMessages() {
+          return messages;
+        },
+        toolCalls,
+      };
+    } else {
+      throw new Error("Cannot get tool calls from response");
+    }
+  }
 }
 
 export abstract class ToolCallLLM<

packages/core/src/llms/mock.ts

@@ -1,5 +1,4 @@
-// TODO: move to a test package
-import { ToolCallLLM } from "../llms/base";
+import { ToolCallLLM } from "./base";
 import type {
   ChatResponse,
   ChatResponseChunk,
@@ -9,7 +8,7 @@ import type {
   LLMCompletionParamsNonStreaming,
   LLMCompletionParamsStreaming,
   LLMMetadata,
-} from "../llms/type";
+} from "./type";
 
 export class MockLLM extends ToolCallLLM {
   metadata: LLMMetadata;

packages/core/src/llms/tool-call.ts

@@ -0,0 +1,61 @@
+import { stringifyJSONToMessageContent } from "../utils";
+import type {
+  BaseTool,
+  ChatMessage,
+  ChatResponse,
+  ChatResponseChunk,
+  ToolCall,
+  ToolCallLLMMessageOptions,
+} from "./type";
+
+export const getToolCallsFromResponse = (
+  response:
+    | ChatResponse<ToolCallLLMMessageOptions>
+    | ChatResponseChunk<ToolCallLLMMessageOptions>,
+): ToolCall[] => {
+  let options;
+
+  if ("message" in response) {
+    options = response.message.options;
+  } else {
+    options = response.options;
+  }
+
+  if (options && "toolCall" in options) {
+    return (options.toolCall as ToolCall[]).map((toolCall) => ({
+      ...toolCall,
+      input:
+        // XXX: this is a hack openai returns parsed object for streaming, but not for
+        // non-streaming
+        typeof toolCall.input === "string"
+          ? JSON.parse(toolCall.input)
+          : toolCall.input,
+    }));
+  }
+  return [];
+};
+
+export const callTool = async <
+  AdditionalMessageOptions extends object = object,
+>(
+  tools: BaseTool[],
+  toolCall: ToolCall,
+): Promise<ChatMessage<AdditionalMessageOptions> | null> => {
+  const tool = tools?.find((t) => t.metadata.name === toolCall.name);
+  // TODO: consider using BaseToolWithCall instead of BaseTool to avoid checking for tool.call
+  if (tool && tool.call) {
+    const result = await tool.call(toolCall.input);
+    const toolResultMessage: ChatMessage<AdditionalMessageOptions> = {
+      role: "user",
+      content: stringifyJSONToMessageContent(result),
+      options: {
+        toolResult: {
+          id: toolCall.id,
+          result,
+        },
+      } as AdditionalMessageOptions,
+    };
+    return toolResultMessage;
+  }
+  return null;
+};

packages/core/src/llms/type.ts

@@ -95,6 +95,22 @@ export type ChatResponseChunk<
   options?: undefined | AdditionalMessageOptions;
 };
 
+export interface ExecResponse<
+  AdditionalMessageOptions extends object = object,
+> {
+  newMessages: ChatMessage<AdditionalMessageOptions>[];
+  toolCalls: ToolCall[];
+}
+
+export interface ExecStreamResponse<
+  AdditionalMessageOptions extends object = object,
+> {
+  stream: AsyncIterable<ChatResponseChunk<AdditionalMessageOptions>>;
+  // this is a function as while streaming, the assistant message is not ready yet - can be called after the stream is done
+  newMessages(): ChatMessage<AdditionalMessageOptions>[];
+  toolCalls: ToolCall[];
+}
+
 export interface CompletionResponse {
   text: string;
   /**
@@ -120,9 +136,9 @@ export interface LLMChatParamsBase<
   AdditionalMessageOptions extends object = object,
 > {
   messages: ChatMessage<AdditionalMessageOptions>[];
-  additionalChatOptions?: AdditionalChatOptions;
-  tools?: BaseTool[];
-  responseFormat?: z.ZodType | object;
+  additionalChatOptions?: AdditionalChatOptions | undefined;
+  tools?: BaseTool[] | undefined;
+  responseFormat?: z.ZodType | object | undefined;
 }
 
 export interface LLMChatParamsStreaming<

packages/core/src/utils/index.ts

@@ -70,8 +70,6 @@ export {
   toToolDescriptions,
 } from "./llms";
 
-export { MockLLM } from "./mock";
-
 export * from "./encoding";
 export { objectEntries } from "./object-entries";
 export * from "./stream";

packages/core/tests/agents.test.ts

@@ -1,5 +1,5 @@
 import { LLMAgent, validateAgentParams } from "@llamaindex/core/agent";
-import { MockLLM } from "@llamaindex/core/utils";
+import { MockLLM } from "@llamaindex/core/llms/mock";
 import { expect, test } from "vitest";
 import { ZodError } from "zod";
 

packages/core/tests/llms/base.test.ts

@@ -0,0 +1,43 @@
+import { MockLLM } from "@llamaindex/core/llms/mock";
+import { describe, expect, it } from "vitest";
+
+// TODO: add tests for tool calls
+describe("BaseLLM exec", () => {
+  it("should stream text response when no tool call is made", async () => {
+    const responseMessage = "This is a response message while streaming";
+
+    const llm = new MockLLM({ responseMessage });
+
+    const { stream, newMessages, toolCalls } = await llm.exec({
+      messages: [{ content: "Hi", role: "user" }],
+      stream: true,
+    });
+
+    expect(() => newMessages()).toThrowError();
+
+    const chunks = [];
+    for await (const chunk of stream) {
+      chunks.push(chunk);
+    }
+
+    expect(chunks.map((c) => c.delta).join("")).toBe(responseMessage);
+    expect(toolCalls).toEqual([]);
+    expect(newMessages()).toEqual([
+      { content: responseMessage, role: "assistant" },
+    ]);
+  });
+  it("should return text response when no tool call is made", async () => {
+    const responseMessage = "This is a response message";
+
+    const llm = new MockLLM({ responseMessage });
+
+    const { newMessages, toolCalls } = await llm.exec({
+      messages: [{ content: "Hi", role: "user" }],
+    });
+
+    expect(newMessages).toEqual([
+      { content: responseMessage, role: "assistant" },
+    ]);
+    expect(toolCalls).toEqual([]);
+  });
+});

packages/core/tests/memory/memory.test.ts

@@ -1,7 +1,7 @@
 import { Settings } from "@llamaindex/core/global";
 import type { ChatMessage, LLM } from "@llamaindex/core/llms";
+import { MockLLM } from "@llamaindex/core/llms/mock";
 import { createMemory, Memory, staticBlock } from "@llamaindex/core/memory";
-import { MockLLM } from "@llamaindex/core/utils";
 import type { Tokenizer } from "@llamaindex/env/tokenizers";
 import {
   afterAll,

packages/core/tests/simple-chat-engine.test.ts

@@ -1,6 +1,6 @@
 import { SimpleChatEngine } from "@llamaindex/core/chat-engine";
+import { MockLLM } from "@llamaindex/core/llms/mock";
 import { Memory } from "@llamaindex/core/memory";
-import { MockLLM } from "@llamaindex/core/utils";
 import { describe, expect, test } from "vitest";
 
 describe("SimpleChatEngine", () => {

packages/tools/src/tool-call.ts

@@ -135,6 +135,9 @@ class ChatWithToolsResponse {
   }
 }
 
+/**
+ * @deprecated Use @BaseLLM.exec instead.
+ */
 export const chatWithTools = async (
   llm: ToolCallLLM,
   tools: BaseToolWithCall[],

packages/tools/tests/code-generator.test.ts

@@ -1,5 +1,5 @@
 import { Settings } from "@llamaindex/core/global";
-import { MockLLM } from "@llamaindex/core/utils";
+import { MockLLM } from "@llamaindex/core/llms/mock";
 import { describe, expect, test } from "vitest";
 import {
   codeGenerator,

packages/workflow/test/agent-workflow.test.ts

@@ -1,6 +1,6 @@
 import { Settings } from "@llamaindex/core/global";
+import { MockLLM } from "@llamaindex/core/llms/mock";
 import { FunctionTool } from "@llamaindex/core/tools";
-import { MockLLM } from "@llamaindex/core/utils";
 import { describe, expect, test, vi } from "vitest";
 import { z } from "zod";
 import {

packages/workflow/test/function-agent.test.ts

@@ -1,6 +1,6 @@
 import { ChatMessage } from "@llamaindex/core/llms";
+import { MockLLM } from "@llamaindex/core/llms/mock";
 import { tool } from "@llamaindex/core/tools";
-import { MockLLM } from "@llamaindex/core/utils";
 import { type WorkflowContext } from "@llamaindex/workflow-core";
 import { zodEvent } from "@llamaindex/workflow-core/util/zod";
 import { describe, expect, test, vi } from "vitest";

packages/workflow/test/mock.ts

@@ -4,7 +4,7 @@ import {
   LLMChatParamsStreaming,
   ToolCallLLM,
 } from "@llamaindex/core/llms";
-import { MockLLM } from "@llamaindex/core/utils";
+import { MockLLM } from "@llamaindex/core/llms/mock";
 import { vi } from "vitest";
 
 /**