From 698cd9c631726efca5b2a9f104f398010be03602 Mon Sep 17 00:00:00 2001
From: Emanuel Ferreira <contatoferreirads@gmail.com>
Date: Fri, 1 Mar 2024 21:28:02 -0300
Subject: [PATCH] fix: step wise agent + examples (#594)
---
.changeset/four-ways-enjoy.md | 5 ++
examples/agent/step_wise_openai.ts | 95 ++++++++++++++++++++++++++
examples/agent/step_wise_query_tool.ts | 64 +++++++++++++++++
examples/agent/step_wise_react.ts | 90 ++++++++++++++++++++++++
packages/core/src/agent/runner/base.ts | 5 +-
packages/core/src/agent/types.ts | 4 +-
packages/core/src/llm/LLM.ts | 2 +-
7 files changed, 260 insertions(+), 5 deletions(-)
create mode 100644 .changeset/four-ways-enjoy.md
create mode 100644 examples/agent/step_wise_openai.ts
create mode 100644 examples/agent/step_wise_query_tool.ts
create mode 100644 examples/agent/step_wise_react.ts
diff --git a/.changeset/four-ways-enjoy.md b/.changeset/four-ways-enjoy.md
new file mode 100644
index 000000000..fee00e66f
--- /dev/null
+++ b/.changeset/four-ways-enjoy.md
@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+fix: step wise agent + examples
diff --git a/examples/agent/step_wise_openai.ts b/examples/agent/step_wise_openai.ts
new file mode 100644
index 000000000..abbc273c5
--- /dev/null
+++ b/examples/agent/step_wise_openai.ts
@@ -0,0 +1,95 @@
+import { FunctionTool, OpenAIAgent } from "llamaindex";
+
+// Define a function to sum two numbers
+function sumNumbers({ a, b }: { a: number; b: number }): number {
+ return a + b;
+}
+
+// Define a function to divide two numbers
+function divideNumbers({ a, b }: { a: number; b: number }): number {
+ return a / b;
+}
+
+// Define the parameters of the sum function as a JSON schema
+const sumJSON = {
+ type: "object",
+ properties: {
+ a: {
+ type: "number",
+ description: "The first number",
+ },
+ b: {
+ type: "number",
+ description: "The second number",
+ },
+ },
+ required: ["a", "b"],
+};
+
+const divideJSON = {
+ type: "object",
+ properties: {
+ a: {
+ type: "number",
+ description: "The dividend a to divide",
+ },
+ b: {
+ type: "number",
+ description: "The divisor b to divide by",
+ },
+ },
+ required: ["a", "b"],
+};
+
+async function main() {
+ // Create a function tool from the sum function
+ const functionTool = new FunctionTool(sumNumbers, {
+ name: "sumNumbers",
+ description: "Use this function to sum two numbers",
+ parameters: sumJSON,
+ });
+
+ // Create a function tool from the divide function
+ const functionTool2 = new FunctionTool(divideNumbers, {
+ name: "divideNumbers",
+ description: "Use this function to divide two numbers",
+ parameters: divideJSON,
+ });
+
+ // Create an OpenAIAgent with the function tools
+ const agent = new OpenAIAgent({
+ tools: [functionTool, functionTool2],
+ verbose: true,
+ });
+
+ // Create a task to sum and divide numbers
+ const task = agent.createTask("How much is 5 + 5? then divide by 2");
+
+ let count = 0;
+
+ while (true) {
+ const stepOutput = await agent.runStep(task.taskId);
+
+ console.log(`Runnning step ${count++}`);
+ console.log(`======== OUTPUT ==========`);
+ if (stepOutput.output.response) {
+ console.log(stepOutput.output.response);
+ } else {
+ console.log(stepOutput.output.sources);
+ }
+ console.log(`==========================`);
+
+ if (stepOutput.isLast) {
+ const finalResponse = await agent.finalizeResponse(
+ task.taskId,
+ stepOutput,
+ );
+ console.log({ finalResponse });
+ break;
+ }
+ }
+}
+
+main().then(() => {
+ console.log("Done");
+});
diff --git a/examples/agent/step_wise_query_tool.ts b/examples/agent/step_wise_query_tool.ts
new file mode 100644
index 000000000..d7a17295a
--- /dev/null
+++ b/examples/agent/step_wise_query_tool.ts
@@ -0,0 +1,64 @@
+import {
+ OpenAIAgent,
+ QueryEngineTool,
+ SimpleDirectoryReader,
+ VectorStoreIndex,
+} from "llamaindex";
+
+async function main() {
+ // Load the documents
+ const documents = await new SimpleDirectoryReader().loadData({
+ directoryPath: "node_modules/llamaindex/examples",
+ });
+
+ // Create a vector index from the documents
+ const vectorIndex = await VectorStoreIndex.fromDocuments(documents);
+
+ // Create a query engine from the vector index
+ const abramovQueryEngine = vectorIndex.asQueryEngine();
+
+ // Create a QueryEngineTool with the query engine
+ const queryEngineTool = new QueryEngineTool({
+ queryEngine: abramovQueryEngine,
+ metadata: {
+ name: "abramov_query_engine",
+ description: "A query engine for the Abramov documents",
+ },
+ });
+
+ // Create an OpenAIAgent with the function tools
+ const agent = new OpenAIAgent({
+ tools: [queryEngineTool],
+ verbose: true,
+ });
+
+ const task = agent.createTask("What was his salary?");
+
+ let count = 0;
+
+ while (true) {
+ const stepOutput = await agent.runStep(task.taskId);
+
+ console.log(`Runnning step ${count++}`);
+ console.log(`======== OUTPUT ==========`);
+ if (stepOutput.output.response) {
+ console.log(stepOutput.output.response);
+ } else {
+ console.log(stepOutput.output.sources);
+ }
+ console.log(`==========================`);
+
+ if (stepOutput.isLast) {
+ const finalResponse = await agent.finalizeResponse(
+ task.taskId,
+ stepOutput,
+ );
+ console.log({ finalResponse });
+ break;
+ }
+ }
+}
+
+main().then(() => {
+ console.log("Done");
+});
diff --git a/examples/agent/step_wise_react.ts b/examples/agent/step_wise_react.ts
new file mode 100644
index 000000000..4230db586
--- /dev/null
+++ b/examples/agent/step_wise_react.ts
@@ -0,0 +1,90 @@
+import { FunctionTool, ReActAgent } from "llamaindex";
+
+// Define a function to sum two numbers
+function sumNumbers({ a, b }: { a: number; b: number }): number {
+ return a + b;
+}
+
+// Define a function to divide two numbers
+function divideNumbers({ a, b }: { a: number; b: number }): number {
+ return a / b;
+}
+
+// Define the parameters of the sum function as a JSON schema
+const sumJSON = {
+ type: "object",
+ properties: {
+ a: {
+ type: "number",
+ description: "The first number",
+ },
+ b: {
+ type: "number",
+ description: "The second number",
+ },
+ },
+ required: ["a", "b"],
+};
+
+const divideJSON = {
+ type: "object",
+ properties: {
+ a: {
+ type: "number",
+ description: "The dividend",
+ },
+ b: {
+ type: "number",
+ description: "The divisor",
+ },
+ },
+ required: ["a", "b"],
+};
+
+async function main() {
+ // Create a function tool from the sum function
+ const functionTool = new FunctionTool(sumNumbers, {
+ name: "sumNumbers",
+ description: "Use this function to sum two numbers",
+ parameters: sumJSON,
+ });
+
+ // Create a function tool from the divide function
+ const functionTool2 = new FunctionTool(divideNumbers, {
+ name: "divideNumbers",
+ description: "Use this function to divide two numbers",
+ parameters: divideJSON,
+ });
+
+ // Create an OpenAIAgent with the function tools
+ const agent = new ReActAgent({
+ tools: [functionTool, functionTool2],
+ verbose: true,
+ });
+
+ const task = agent.createTask("Divide 16 by 2 then add 20");
+
+ let count = 0;
+
+ while (true) {
+ const stepOutput = await agent.runStep(task.taskId);
+
+ console.log(`Runnning step ${count++}`);
+ console.log(`======== OUTPUT ==========`);
+ console.log(stepOutput.output);
+ console.log(`==========================`);
+
+ if (stepOutput.isLast) {
+ const finalResponse = await agent.finalizeResponse(
+ task.taskId,
+ stepOutput,
+ );
+ console.log({ finalResponse });
+ break;
+ }
+ }
+}
+
+main().then(() => {
+ console.log("Done");
+});
diff --git a/packages/core/src/agent/runner/base.ts b/packages/core/src/agent/runner/base.ts
index 80dfd2f71..532390e0d 100644
--- a/packages/core/src/agent/runner/base.ts
+++ b/packages/core/src/agent/runner/base.ts
@@ -14,7 +14,7 @@ import { AgentState, BaseAgentRunner, TaskState } from "./types.js";
const validateStepFromArgs = (
taskId: string,
- input: string,
+ input?: string | null,
step?: any,
kwargs?: any,
): TaskStep | undefined => {
@@ -24,6 +24,7 @@ const validateStepFromArgs = (
}
return step;
} else {
+ if (!input) return;
return new TaskStep(taskId, step, input, kwargs);
}
};
@@ -194,7 +195,7 @@ export class AgentRunner extends BaseAgentRunner {
*/
async runStep(
taskId: string,
- input: string,
+ input?: string | null,
step?: TaskStep,
kwargs: any = {},
): Promise<TaskStepOutput> {
diff --git a/packages/core/src/agent/types.ts b/packages/core/src/agent/types.ts
index f861a8488..8adc96340 100644
--- a/packages/core/src/agent/types.ts
+++ b/packages/core/src/agent/types.ts
@@ -161,13 +161,13 @@ export class TaskStep implements ITaskStep {
* @param isLast: isLast
*/
export class TaskStepOutput {
- output: unknown;
+ output: any;
taskStep: TaskStep;
nextSteps: TaskStep[];
isLast: boolean;
constructor(
- output: unknown,
+ output: any,
taskStep: TaskStep,
nextSteps: TaskStep[],
isLast: boolean = false,
diff --git a/packages/core/src/llm/LLM.ts b/packages/core/src/llm/LLM.ts
index 317fcc1a3..4db929e79 100644
--- a/packages/core/src/llm/LLM.ts
+++ b/packages/core/src/llm/LLM.ts
@@ -260,7 +260,7 @@ export class OpenAI extends BaseLLM {
stream: false,
});
- const content = response.choices[0].message?.content ?? "";
+ const content = response.choices[0].message?.content ?? null;
const kwargsOutput: Record<string, any> = {};
--
GitLab