langchain[minor]: feat(embedding integration): Gradient AI

docs/api_refs/typedoc.json

@@ -67,6 +67,7 @@
     "./langchain/src/embeddings/minimax.ts",
     "./langchain/src/embeddings/voyage.ts",
     "./langchain/src/embeddings/llama_cpp.ts",
+    "./langchain/src/embeddings/gradient_ai.ts",
     "./langchain/src/llms/load.ts",
     "./langchain/src/llms/base.ts",
     "./langchain/src/llms/openai.ts",

docs/core_docs/docs/integrations/text_embedding/gradient_ai.mdx

@@ -0,0 +1,39 @@
+---
+sidebar_class_name: node-only
+---
+
+# Gradient AI
+
+The `GradientEmbeddings` class uses the Gradient AI API to generate embeddings for a given text.
+
+## Setup
+
+You'll need to install the official Gradient Node SDK as a peer dependency:
+
+```bash npm2yarn
+npm i @gradientai/nodejs-sdk
+```
+
+You will need to set the following environment variables for using the Gradient AI API.
+
+```
+export GRADIENT_ACCESS_TOKEN=<YOUR_ACCESS_TOKEN>
+export GRADIENT_WORKSPACE_ID=<YOUR_WORKSPACE_ID>
+```
+
+Alternatively, these can be set during the GradientAI Class instantiation as `gradientAccessKey` and `workspaceId` respectively.
+For example:
+
+```typescript
+const model = new GradientEmbeddings({
+  gradientAccessKey: "My secret Access Token"
+  workspaceId: "My secret workspace id"
+});
+```
+
+## Usage
+
+import CodeBlock from "@theme/CodeBlock";
+import GradientEmbeddingsExample from "@examples/embeddings/gradient_ai.ts";
+
+<CodeBlock language="typescript">{GradientEmbeddingsExample}</CodeBlock>

examples/src/embeddings/gradient_ai.ts

@@ -0,0 +1,7 @@
+import { GradientEmbeddings } from "langchain/embeddings/gradient_ai";
+
+const model = new GradientEmbeddings();
+const res = await model.embedQuery(
+  "What would be a good company name a company that makes colorful socks?"
+);
+console.log({ res });

langchain/.gitignore

@@ -145,6 +145,9 @@ embeddings/voyage.d.ts
 embeddings/llama_cpp.cjs
 embeddings/llama_cpp.js
 embeddings/llama_cpp.d.ts
+embeddings/gradient_ai.cjs
+embeddings/gradient_ai.js
+embeddings/gradient_ai.d.ts
 llms/load.cjs
 llms/load.js
 llms/load.d.ts

langchain/package.json

@@ -157,6 +157,9 @@
     "embeddings/llama_cpp.cjs",
     "embeddings/llama_cpp.js",
     "embeddings/llama_cpp.d.ts",
+    "embeddings/gradient_ai.cjs",
+    "embeddings/gradient_ai.js",
+    "embeddings/gradient_ai.d.ts",
     "llms/load.cjs",
     "llms/load.js",
     "llms/load.d.ts",
@@ -1698,6 +1701,11 @@
       "import": "./embeddings/llama_cpp.js",
       "require": "./embeddings/llama_cpp.cjs"
     },
+    "./embeddings/gradient_ai": {
+      "types": "./embeddings/gradient_ai.d.ts",
+      "import": "./embeddings/gradient_ai.js",
+      "require": "./embeddings/gradient_ai.cjs"
+    },
     "./llms/load": {
       "types": "./llms/load.d.ts",
       "import": "./llms/load.js",

langchain/scripts/create-entrypoints.js

@@ -63,6 +63,7 @@ const entrypoints = {
   "embeddings/minimax": "embeddings/minimax",
   "embeddings/voyage": "embeddings/voyage",
   "embeddings/llama_cpp": "embeddings/llama_cpp",
+  "embeddings/gradient_ai": "embeddings/gradient_ai",
   // llms
   "llms/load": "llms/load",
   "llms/base": "llms/base",
@@ -364,6 +365,7 @@ const requiresOptionalDependency = [
   "embeddings/hf",
   "embeddings/hf_transformers",
   "embeddings/llama_cpp",
+  "embeddings/gradient_ai",
   "llms/load",
   "llms/cohere",
   "llms/googlevertexai",

langchain/src/embeddings/gradient_ai.ts

@@ -0,0 +1,117 @@
+import { Gradient } from "@gradientai/nodejs-sdk";
+import { getEnvironmentVariable } from "../util/env.js";
+import { chunkArray } from "../util/chunk.js";
+import { Embeddings, EmbeddingsParams } from "./base.js";
+
+/**
+ * Interface for GradientEmbeddings parameters. Extends EmbeddingsParams and
+ * defines additional parameters specific to the GradientEmbeddings class.
+ */
+export interface GradientEmbeddingsParams extends EmbeddingsParams {
+  /**
+   * Gradient AI Access Token.
+   * Provide Access Token if you do not wish to automatically pull from env.
+   */
+  gradientAccessKey?: string;
+  /**
+   * Gradient Workspace Id.
+   * Provide workspace id if you do not wish to automatically pull from env.
+   */
+  workspaceId?: string;
+}
+
+/**
+ * Class for generating embeddings using the Gradient AI's API. Extends the
+ * Embeddings class and implements GradientEmbeddingsParams and
+ */
+export class GradientEmbeddings
+  extends Embeddings
+  implements GradientEmbeddingsParams
+{
+  gradientAccessKey?: string;
+
+  workspaceId?: string;
+
+  batchSize = 128;
+
+  model: any;
+
+  constructor(fields: GradientEmbeddingsParams) {
+    super(fields);
+
+    this.gradientAccessKey =
+      fields?.gradientAccessKey ??
+      getEnvironmentVariable("GRADIENT_ACCESS_TOKEN");
+    this.workspaceId =
+      fields?.workspaceId ?? getEnvironmentVariable("GRADIENT_WORKSPACE_ID");
+
+    if (!this.gradientAccessKey) {
+      throw new Error("Missing Gradient AI Access Token");
+    }
+
+    if (!this.workspaceId) {
+      throw new Error("Missing Gradient AI Workspace ID");
+    }
+  }
+
+  /**
+   * Method to generate embeddings for an array of documents. Splits the
+   * documents into batches and makes requests to the Gradient API to generate
+   * embeddings.
+   * @param texts Array of documents to generate embeddings for.
+   * @returns Promise that resolves to a 2D array of embeddings for each document.
+   */
+  async embedDocuments(texts: string[]): Promise<number[][]> {
+    await this.setModel();
+
+    const mappedTexts = texts.map((text) => ({ input: text }));
+
+    const batches = chunkArray(mappedTexts, this.batchSize);
+
+    const batchRequests = batches.map((batch) =>
+      this.caller.call(async () =>
+        this.model.generateEmbeddings({
+          inputs: batch,
+        })
+      )
+    );
+    const batchResponses = await Promise.all(batchRequests);
+
+    const embeddings: number[][] = [];
+    for (let i = 0; i < batchResponses.length; i += 1) {
+      const batch = batches[i];
+      const { embeddings: batchResponse } = batchResponses[i];
+      for (let j = 0; j < batch.length; j += 1) {
+        embeddings.push(batchResponse[j].embedding);
+      }
+    }
+    return embeddings;
+  }
+
+  /**
+   * Method to generate an embedding for a single document. Calls the
+   * embedDocuments method with the document as the input.
+   * @param text Document to generate an embedding for.
+   * @returns Promise that resolves to an embedding for the document.
+   */
+  async embedQuery(text: string): Promise<number[]> {
+    const data = await this.embedDocuments([text]);
+    return data[0];
+  }
+
+  /**
+   * Method to set the model to use for generating embeddings.
+   * @sets the class' `model` value to that of the retrieved Embeddings Model.
+   */
+  async setModel() {
+    if (this.model) return;
+
+    const gradient = new Gradient({
+      accessToken: this.gradientAccessKey,
+      workspaceId: this.workspaceId,
+    });
+    this.model = await gradient.getEmbeddingsModel({
+      slug: "bge-large",
+    });
+  }
+}

langchain/src/load/import_constants.ts

@@ -24,6 +24,7 @@ export const optionalImportEntrypoints = [
   "langchain/embeddings/googlevertexai",
   "langchain/embeddings/googlepalm",
   "langchain/embeddings/llama_cpp",
+  "langchain/embeddings/gradient_ai",
   "langchain/llms/load",
   "langchain/llms/cohere",
   "langchain/llms/hf",

langchain/src/load/import_type.d.ts

@@ -70,6 +70,9 @@ export interface OptionalImportMap {
   "langchain/embeddings/llama_cpp"?:
     | typeof import("../embeddings/llama_cpp.js")
     | Promise<typeof import("../embeddings/llama_cpp.js")>;
+  "langchain/embeddings/gradient_ai"?:
+    | typeof import("../embeddings/gradient_ai.js")
+    | Promise<typeof import("../embeddings/gradient_ai.js")>;
   "langchain/llms/load"?:
     | typeof import("../llms/load.js")
     | Promise<typeof import("../llms/load.js")>;