From e37fa5d9ca0a570451f72a33654edc43af3fb713 Mon Sep 17 00:00:00 2001
From: Marcus Schiesser <mail@marcusschiesser.de>
Date: Tue, 7 May 2024 10:41:14 +0800
Subject: [PATCH] docs: add retriever tool example (#814)

---
 examples/agent/query_openai_agent.ts     |  5 +-
 examples/agent/retriever_openai_agent.ts | 65 ++++++++++++++++++++++++
 2 files changed, 68 insertions(+), 2 deletions(-)
 create mode 100644 examples/agent/retriever_openai_agent.ts

diff --git a/examples/agent/query_openai_agent.ts b/examples/agent/query_openai_agent.ts
index 4521c9921c..57dc838f59 100644
--- a/examples/agent/query_openai_agent.ts
+++ b/examples/agent/query_openai_agent.ts
@@ -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(() => {
diff --git a/examples/agent/retriever_openai_agent.ts b/examples/agent/retriever_openai_agent.ts
new file mode 100644
index 0000000000..09e940951b
--- /dev/null
+++ b/examples/agent/retriever_openai_agent.ts
@@ -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");
+});