feat: support OpenaiAssistantAgent

.changeset/yellow-grapes-explain.md

@@ -0,0 +1,6 @@
+---
+"llamaindex": patch
+"llamaindex-loader-example": patch
+---
+
+feat: support `OpenaiAssistantAgent`

.npmrc

@@ -2,3 +2,4 @@ auto-install-peers = true
 enable-pre-post-scripts = true
 prefer-workspace-packages = true
 save-workspace-protocol = true
+link-workspace-packages = true

examples/agent/openai-assistant-agent.ts

@@ -0,0 +1,13 @@
+import { OpenaiAssistantAgent } from 'llamaindex/agent/openai'
+
+const assistantId = process.env.ASSISTANT_ID
+
+if (!assistantId) {
+  throw new Error('ASSISTANT_ID is required for openai')
+}
+
+const agent = new OpenaiAssistantAgent({
+  assistantId
+})
+
+agent.run('What\'s the weather today?')

examples/package.json

@@ -12,7 +12,7 @@
     "commander": "^11.1.0",
     "dotenv": "^16.4.5",
     "js-tiktoken": "^1.0.11",
-    "llamaindex": "*",
+    "llamaindex": "latest",
     "mongodb": "^6.5.0",
     "pathe": "^1.1.2"
   },

examples/readers/package.json

@@ -12,7 +12,7 @@
     "start:llamaparse": "node --loader ts-node/esm ./src/llamaparse.ts"
   },
   "dependencies": {
-    "llamaindex": "*"
+    "llamaindex": "latest"
   },
   "devDependencies": {
     "@types/node": "^20.12.7",

packages/core/src/agent/index.ts

@@ -7,6 +7,8 @@ export {
   OpenAIAgent,
   OpenAIAgentWorker,
   type OpenAIAgentParams,
+  OpenaiAssistantAgent,
+  type OpenaiAssistantAgentParams
 } from "./openai.js";
 export {
   ReACTAgent,

packages/core/src/agent/openai.ts

@@ -1,194 +1,2 @@
-import { pipeline } from "@llamaindex/env";
-import { Settings } from "../Settings.js";
-import type {
-  ChatResponseChunk,
-  ToolCall,
-  ToolCallLLMMessageOptions,
-} from "../llm/index.js";
-import { OpenAI } from "../llm/open_ai.js";
-import { ObjectRetriever } from "../objects/index.js";
-import type { BaseToolWithCall } from "../types.js";
-import {
-  AgentRunner,
-  AgentWorker,
-  type AgentParamsBase,
-  type TaskHandler,
-} from "./base.js";
-import { callTool } from "./utils.js";
-
-type OpenAIParamsBase = AgentParamsBase<OpenAI>;
-
-type OpenAIParamsWithTools = OpenAIParamsBase & {
-  tools: BaseToolWithCall[];
-};
-
-type OpenAIParamsWithToolRetriever = OpenAIParamsBase & {
-  toolRetriever: ObjectRetriever<BaseToolWithCall>;
-};
-
-export type OpenAIAgentParams =
-  | OpenAIParamsWithTools
-  | OpenAIParamsWithToolRetriever;
-
-export class OpenAIAgentWorker extends AgentWorker<OpenAI> {
-  taskHandler = OpenAIAgent.taskHandler;
-}
-
-export class OpenAIAgent extends AgentRunner<OpenAI> {
-  constructor(params: OpenAIAgentParams) {
-    super({
-      llm:
-        params.llm ?? Settings.llm instanceof OpenAI
-          ? (Settings.llm as OpenAI)
-          : new OpenAI(),
-      chatHistory: params.chatHistory ?? [],
-      runner: new OpenAIAgentWorker(),
-      systemPrompt: params.systemPrompt ?? null,
-      tools:
-        "tools" in params
-          ? params.tools
-          : params.toolRetriever.retrieve.bind(params.toolRetriever),
-    });
-  }
-
-  createStore = AgentRunner.defaultCreateStore;
-
-  static taskHandler: TaskHandler<OpenAI> = async (step) => {
-    const { input } = step;
-    const { llm, stream, getTools } = step.context;
-    if (input) {
-      step.context.store.messages = [...step.context.store.messages, input];
-    }
-    const lastMessage = step.context.store.messages.at(-1)!.content;
-    const tools = await getTools(lastMessage);
-    const response = await llm.chat({
-      // @ts-expect-error
-      stream,
-      tools,
-      messages: [...step.context.store.messages],
-    });
-    if (!stream) {
-      step.context.store.messages = [
-        ...step.context.store.messages,
-        response.message,
-      ];
-      const options = response.message.options ?? {};
-      if ("toolCall" in options) {
-        const { toolCall } = options;
-        const targetTool = tools.find(
-          (tool) => tool.metadata.name === toolCall.name,
-        );
-        const toolOutput = await callTool(targetTool, toolCall);
-        step.context.store.toolOutputs.push(toolOutput);
-        return {
-          taskStep: step,
-          output: {
-            raw: response.raw,
-            message: {
-              content: toolOutput.output,
-              role: "user",
-              options: {
-                toolResult: {
-                  result: toolOutput.output,
-                  isError: toolOutput.isError,
-                  id: toolCall.id,
-                },
-              },
-            },
-          },
-          isLast: false,
-        };
-      } else {
-        return {
-          taskStep: step,
-          output: response,
-          isLast: true,
-        };
-      }
-    } else {
-      const responseChunkStream = new ReadableStream<
-        ChatResponseChunk<ToolCallLLMMessageOptions>
-      >({
-        async start(controller) {
-          for await (const chunk of response) {
-            controller.enqueue(chunk);
-          }
-          controller.close();
-        },
-      });
-      const [pipStream, finalStream] = responseChunkStream.tee();
-      const reader = pipStream.getReader();
-      const { value } = await reader.read();
-      reader.releaseLock();
-      if (value === undefined) {
-        throw new Error(
-          "first chunk value is undefined, this should not happen",
-        );
-      }
-      // check if first chunk has tool calls, if so, this is a function call
-      // otherwise, it's a regular message
-      const hasToolCall = !!(value.options && "toolCall" in value.options);
-
-      if (hasToolCall) {
-        // you need to consume the response to get the full toolCalls
-        const toolCalls = await pipeline(
-          pipStream,
-          async (
-            iter: AsyncIterable<ChatResponseChunk<ToolCallLLMMessageOptions>>,
-          ) => {
-            const toolCalls = new Map<string, ToolCall>();
-            for await (const chunk of iter) {
-              if (chunk.options && "toolCall" in chunk.options) {
-                const toolCall = chunk.options.toolCall;
-                toolCalls.set(toolCall.id, toolCall);
-              }
-            }
-            return [...toolCalls.values()];
-          },
-        );
-        for (const toolCall of toolCalls) {
-          const targetTool = tools.find(
-            (tool) => tool.metadata.name === toolCall.name,
-          );
-          step.context.store.messages = [
-            ...step.context.store.messages,
-            {
-              role: "assistant" as const,
-              content: "",
-              options: {
-                toolCall,
-              },
-            },
-          ];
-          const toolOutput = await callTool(targetTool, toolCall);
-          step.context.store.messages = [
-            ...step.context.store.messages,
-            {
-              role: "user" as const,
-              content: toolOutput.output,
-              options: {
-                toolResult: {
-                  result: toolOutput.output,
-                  isError: toolOutput.isError,
-                  id: toolCall.id,
-                },
-              },
-            },
-          ];
-          step.context.store.toolOutputs.push(toolOutput);
-        }
-        return {
-          taskStep: step,
-          output: null,
-          isLast: false,
-        };
-      } else {
-        return {
-          taskStep: step,
-          output: finalStream,
-          isLast: true,
-        };
-      }
-    }
-  };
-}
+export { OpenAIAgent, OpenAIAgentWorker, type OpenAIAgentParams } from './openai/openai-agent.js'
+export { OpenaiAssistantAgent, type OpenaiAssistantAgentParams } from './openai/openai-assistant-agent.js'