Skip to content

Commit

Permalink
fix: step wise agent + examples (run-llama#594)
Browse files Browse the repository at this point in the history
  • Loading branch information
EmanuelCampos authored Mar 2, 2024
1 parent c744a99 commit 698cd9c
Show file tree
Hide file tree
Showing 7 changed files with 260 additions and 5 deletions.
5 changes: 5 additions & 0 deletions .changeset/four-ways-enjoy.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
"llamaindex": patch
---

fix: step wise agent + examples
95 changes: 95 additions & 0 deletions examples/agent/step_wise_openai.ts
Original file line number Diff line number Diff line change
@@ -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");
});
64 changes: 64 additions & 0 deletions examples/agent/step_wise_query_tool.ts
Original file line number Diff line number Diff line change
@@ -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");
});
90 changes: 90 additions & 0 deletions examples/agent/step_wise_react.ts
Original file line number Diff line number Diff line change
@@ -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");
});
5 changes: 3 additions & 2 deletions packages/core/src/agent/runner/base.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 => {
Expand All @@ -24,6 +24,7 @@ const validateStepFromArgs = (
}
return step;
} else {
if (!input) return;
return new TaskStep(taskId, step, input, kwargs);
}
};
Expand Down Expand Up @@ -194,7 +195,7 @@ export class AgentRunner extends BaseAgentRunner {
*/
async runStep(
taskId: string,
input: string,
input?: string | null,
step?: TaskStep,
kwargs: any = {},
): Promise<TaskStepOutput> {
Expand Down
4 changes: 2 additions & 2 deletions packages/core/src/agent/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
2 changes: 1 addition & 1 deletion packages/core/src/llm/LLM.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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> = {};

Expand Down

0 comments on commit 698cd9c

Please sign in to comment.