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'