Add MetalRetriever (#613)

nfcampos · web-flow · commit 047691f0f757 · 2023-04-06T11:25:05.000+01:00
* Add MetalRetriever

* Add docs and example

* Lint
diff --git a/docs/docs/modules/indexes/retrievers/metal-retriever.mdx b/docs/docs/modules/indexes/retrievers/metal-retriever.mdx
@@ -0,0 +1,16 @@
+# Remote Retriever
+
+This example shows how to use the Metal Retriever in a `RetrievalQAChain` to retrieve documents from Metal index.
+
+## Setup
+
+```bash npm2yarn
+npm i @getmetal/metal-sdk
+```
+
+## Usage
+
+import CodeBlock from "@theme/CodeBlock";
+import Example from "@examples/retrievers/metal.ts";
+
+<CodeBlock language="typescript">{Example}</CodeBlock>
diff --git a/docs/docs/modules/indexes/retrievers/remote-retriever.mdx b/docs/docs/modules/indexes/retrievers/remote-retriever.mdx
@@ -1,6 +1,6 @@
 # Remote Retriever
 
-This example shows how to use the a Remote Retriever in a `RetrievalQAChain` to retrieve documents from a remote server.
+This example shows how to use a Remote Retriever in a `RetrievalQAChain` to retrieve documents from a remote server.
 
 ## Usage
 
diff --git a/docs/docs/modules/indexes/retrievers/supabase-hybrid.mdx b/docs/docs/modules/indexes/retrievers/supabase-hybrid.mdx
diff --git a/examples/package.json b/examples/package.json
@@ -21,6 +21,7 @@
   "license": "MIT",
   "dependencies": {
     "@dqbd/tiktoken": "^1.0.2",
+    "@getmetal/metal-sdk": "^1.0.12",
     "@pinecone-database/pinecone": "^0.0.10",
     "@supabase/supabase-js": "^2.10.0",
     "chromadb": "^1.3.0",
diff --git a/examples/src/retrievers/metal.ts b/examples/src/retrievers/metal.ts
@@ -0,0 +1,17 @@
+import Metal from "@getmetal/metal-sdk";
+import { MetalRetriever } from "langchain/retrievers";
+
+export const run = async () => {
+  const MetalSDK = Metal.default;
+
+  const client = new MetalSDK(
+    process.env.METAL_API_KEY!,
+    process.env.METAL_CLIENT_ID!,
+    process.env.METAL_APP_ID
+  );
+  const retriever = new MetalRetriever({ client });
+
+  const docs = await retriever.getRelevantDocuments("hello");
+
+  console.log(docs);
+};
diff --git a/langchain/package.json b/langchain/package.json
@@ -100,6 +100,7 @@
     "@babel/preset-env": "^7.20.2",
     "@dqbd/tiktoken": "^1.0.2",
     "@faker-js/faker": "^7.6.0",
+    "@getmetal/metal-sdk": "^1.0.12",
     "@huggingface/inference": "^1.5.1",
     "@jest/globals": "^29.5.0",
     "@pinecone-database/pinecone": "^0.0.10",
@@ -147,6 +148,7 @@
   },
   "peerDependencies": {
     "@dqbd/tiktoken": "^1.0.2",
+    "@getmetal/metal-sdk": "*",
     "@huggingface/inference": "^1.5.1",
     "@pinecone-database/pinecone": "^0.0.10",
     "@supabase/supabase-js": "^2.10.0",
@@ -169,6 +171,9 @@
     "@dqbd/tiktoken": {
       "optional": true
     },
+    "@getmetal/metal-sdk": {
+      "optional": true
+    },
     "@huggingface/inference": {
       "optional": true
     },
diff --git a/langchain/src/retrievers/index.ts b/langchain/src/retrievers/index.ts
@@ -5,3 +5,4 @@ export {
   SupabaseHybridSearchParams,
 } from "./supabase-hybrid-search.js";
 export { RemoteLangChainRetriever } from "./remote/remote-retriever.js";
+export { MetalRetriever } from "./metal.js";
diff --git a/langchain/src/retrievers/metal.ts b/langchain/src/retrievers/metal.ts
@@ -0,0 +1,36 @@
+import type MetalSDK from "@getmetal/metal-sdk";
+
+import { BaseRetriever } from "../schema/index.js";
+import { Document } from "../document.js";
+
+export interface MetalRetrieverFields {
+  client: MetalSDK.default;
+}
+
+interface ResponseItem {
+  text: string;
+  [key: string]: unknown;
+}
+
+export class MetalRetriever extends BaseRetriever {
+  client: MetalSDK.default;
+
+  constructor(fields: MetalRetrieverFields) {
+    super();
+
+    this.client = fields.client;
+  }
+
+  async getRelevantDocuments(query: string): Promise<Document[]> {
+    const res = await this.client.search({ text: query });
+
+    const items = ("data" in res ? res.data : res) as ResponseItem[];
+    return items.map(
+      ({ text, metadata }) =>
+        new Document({
+          pageContent: text,
+          metadata: metadata as Record<string, unknown>,
+        })
+    );
+  }
+}
diff --git a/langchain/src/retrievers/tests/metal.int.test.ts b/langchain/src/retrievers/tests/metal.int.test.ts
@@ -0,0 +1,22 @@
+/* eslint-disable no-process-env */
+/* eslint-disable @typescript-eslint/no-non-null-assertion */
+import { test, expect } from "@jest/globals";
+import Metal from "@getmetal/metal-sdk";
+
+import { MetalRetriever } from "../metal.js";
+
+test("MetalRetriever", async () => {
+  const MetalSDK = Metal.default;
+  const client = new MetalSDK(
+    process.env.METAL_API_KEY!,
+    process.env.METAL_CLIENT_ID!,
+    process.env.METAL_APP_ID
+  );
+  const retriever = new MetalRetriever({ client });
+
+  const docs = await retriever.getRelevantDocuments("hello");
+
+  expect(docs.length).toBeGreaterThan(0);
+
+  console.log(docs);
+});
diff --git a/yarn.lock b/yarn.lock
@@ -2506,6 +2506,15 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@getmetal/metal-sdk@npm:^1.0.12":
+  version: 1.0.12
+  resolution: "@getmetal/metal-sdk@npm:1.0.12"
+  dependencies:
+    axios: ^1.3.2
+  checksum: 46e3f7876b2aaad759c32a348df078c093fae89205e7050040f80300d04a0c7f1973562d1b9df74b453e2aad316fbfa5d88565e38b33c285d474cdb0356f18f3
+  languageName: node
+  linkType: hard
+
 "@hapi/hoek@npm:^9.0.0":
   version: 9.3.0
   resolution: "@hapi/hoek@npm:9.3.0"
@@ -5052,6 +5061,17 @@ __metadata:
   languageName: node
   linkType: hard
 
+"axios@npm:^1.3.2":
+  version: 1.3.4
+  resolution: "axios@npm:1.3.4"
+  dependencies:
+    follow-redirects: ^1.15.0
+    form-data: ^4.0.0
+    proxy-from-env: ^1.1.0
+  checksum: 7440edefcf8498bc3cdf39de00443e8101f249972c83b739c6e880d9d669fea9486372dbe8739e88b3bf8bb1ad15f6106693f206f078f4516fe8fd47b1c3093c
+  languageName: node
+  linkType: hard
+
 "axobject-query@npm:^3.1.1":
   version: 3.1.1
   resolution: "axobject-query@npm:3.1.1"
@@ -8692,7 +8712,7 @@ __metadata:
   languageName: node
   linkType: hard
 
-"follow-redirects@npm:^1.0.0, follow-redirects@npm:^1.14.7, follow-redirects@npm:^1.14.8":
+"follow-redirects@npm:^1.0.0, follow-redirects@npm:^1.14.7, follow-redirects@npm:^1.14.8, follow-redirects@npm:^1.15.0":
   version: 1.15.2
   resolution: "follow-redirects@npm:1.15.2"
   peerDependenciesMeta:
@@ -11387,6 +11407,7 @@ __metadata:
   resolution: "langchain-examples@workspace:examples"
   dependencies:
     "@dqbd/tiktoken": ^1.0.2
+    "@getmetal/metal-sdk": ^1.0.12
     "@pinecone-database/pinecone": ^0.0.10
     "@supabase/supabase-js": ^2.10.0
     "@tsconfig/recommended": ^1.0.2
@@ -11421,6 +11442,7 @@ __metadata:
     "@dqbd/tiktoken": ^1.0.2
     "@faker-js/faker": ^7.6.0
     "@fortaine/fetch-event-source": ^3.0.6
+    "@getmetal/metal-sdk": ^1.0.12
     "@huggingface/inference": ^1.5.1
     "@jest/globals": ^29.5.0
     "@pinecone-database/pinecone": ^0.0.10
@@ -11479,6 +11501,7 @@ __metadata:
     zod: ^3.21.4
   peerDependencies:
     "@dqbd/tiktoken": ^1.0.2
+    "@getmetal/metal-sdk": "*"
     "@huggingface/inference": ^1.5.1
     "@pinecone-database/pinecone": ^0.0.10
     "@supabase/supabase-js": ^2.10.0
@@ -11499,6 +11522,8 @@ __metadata:
   peerDependenciesMeta:
     "@dqbd/tiktoken":
       optional: true
+    "@getmetal/metal-sdk":
+      optional: true
     "@huggingface/inference":
       optional: true
     "@pinecone-database/pinecone":
@@ -14237,7 +14262,7 @@ __metadata:
   languageName: node
   linkType: hard
 
-"proxy-from-env@npm:1.1.0, proxy-from-env@npm:^1.0.0":
+"proxy-from-env@npm:1.1.0, proxy-from-env@npm:^1.0.0, proxy-from-env@npm:^1.1.0":
   version: 1.1.0
   resolution: "proxy-from-env@npm:1.1.0"
   checksum: ed7fcc2ba0a33404958e34d95d18638249a68c430e30fcb6c478497d72739ba64ce9810a24f53a7d921d0c065e5b78e3822759800698167256b04659366ca4d4
