Skip to content

Commit

Permalink
docs: add retriever tool example (#814)
Browse files Browse the repository at this point in the history
  • Loading branch information
marcusschiesser authored May 7, 2024
1 parent 97e4ecd commit e37fa5d
Show file tree
Hide file tree
Showing 2 changed files with 68 additions and 2 deletions.
5 changes: 3 additions & 2 deletions examples/agent/query_openai_agent.ts
Original file line number Diff line number Diff line change
Expand Up @@ -29,15 +29,16 @@ async function main() {
// Create an OpenAIAgent with the function tools
const agent = new OpenAIAgent({
tools: [queryEngineTool],
verbose: true,
});

// Chat with the agent
const response = await agent.chat({
message: "What was his salary?",
message: "What was his first salary?",
});

// Print the response
console.log(String(response));
console.log(response.response);
}

void main().then(() => {
Expand Down
65 changes: 65 additions & 0 deletions examples/agent/retriever_openai_agent.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
import {
FunctionTool,
MetadataMode,
NodeWithScore,
OpenAIAgent,
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);

const retriever = vectorIndex.asRetriever({ similarityTopK: 3 });

const retrieverTool = FunctionTool.from(
async ({ query }: { query: string }) => {
const nodesWithScores = await retriever.retrieve({
query,
});
return nodesWithScores
.map((nodeWithScore: NodeWithScore) =>
nodeWithScore.node.getContent(MetadataMode.NONE),
)
.join("\n");
},
{
name: "get_abramov_info",
description: "Get information about the Abramov documents",
parameters: {
type: "object",
properties: {
query: {
type: "string",
description: "The query about Abramov",
},
},
required: ["query"],
},
},
);

// Create an OpenAIAgent with the function tools
const agent = new OpenAIAgent({
tools: [retrieverTool],
verbose: true,
});

// Chat with the agent
const response = await agent.chat({
message: "What was his first salary?",
});

// Print the response
console.log(response.response);
}

void main().then(() => {
console.log("Done");
});

0 comments on commit e37fa5d

Please sign in to comment.