community[minor]: DeepInfra embeddings integration #1

examples/src/embeddings/deepinfra.ts

@@ -0,0 +1,14 @@
+import { DeepInfraEmbeddings } from "@langchain/community/embeddings/deepinfra";
+
+
+const model = new DeepInfraEmbeddings({
+  apiToken: process.env.DEEPINFRA_API_TOKEN!,
+  batchSize: 1024, // Default value
+  modelName: "sentence-transformers/clip-ViT-B-32", // Default value
+});
+
+const embeddings = await model.embedQuery(
+  "Tell me a story about a dragon and a princess."
+);
+console.log(embeddings);
+

libs/langchain-community/.gitignore

@@ -138,6 +138,10 @@ embeddings/cohere.cjs
 embeddings/cohere.js
 embeddings/cohere.d.ts
 embeddings/cohere.d.cts
+embeddings/deepinfra.cjs
+embeddings/deepinfra.js
+embeddings/deepinfra.d.ts
+embeddings/deepinfra.d.cts
 embeddings/fireworks.cjs
 embeddings/fireworks.js
 embeddings/fireworks.d.ts

libs/langchain-community/langchain.config.js

@@ -64,6 +64,7 @@ export const config = {
     "embeddings/bedrock": "embeddings/bedrock",
     "embeddings/cloudflare_workersai": "embeddings/cloudflare_workersai",
     "embeddings/cohere": "embeddings/cohere",
+    "embeddings/deepinfra": "embeddings/deepinfra",
     "embeddings/fireworks": "embeddings/fireworks",
     "embeddings/googlepalm": "embeddings/googlepalm",
     "embeddings/googlevertexai": "embeddings/googlevertexai",

libs/langchain-community/package.json

@@ -873,6 +873,15 @@
       "import": "./embeddings/cohere.js",
       "require": "./embeddings/cohere.cjs"
     },
+    "./embeddings/deepinfra": {
+      "types": {
+        "import": "./embeddings/deepinfra.d.ts",
+        "require": "./embeddings/deepinfra.d.cts",
+        "default": "./embeddings/deepinfra.d.ts"
+      },
+      "import": "./embeddings/deepinfra.js",
+      "require": "./embeddings/deepinfra.cjs"
+    },
     "./embeddings/fireworks": {
       "types": {
         "import": "./embeddings/fireworks.d.ts",
@@ -2448,6 +2457,10 @@
     "embeddings/cohere.js",
     "embeddings/cohere.d.ts",
     "embeddings/cohere.d.cts",
+    "embeddings/deepinfra.cjs",
+    "embeddings/deepinfra.js",
+    "embeddings/deepinfra.d.ts",
+    "embeddings/deepinfra.d.cts",
     "embeddings/fireworks.cjs",
     "embeddings/fireworks.js",
     "embeddings/fireworks.d.ts",

libs/langchain-community/src/embeddings/deepinfra.ts

@@ -0,0 +1,210 @@
+import axios, {AxiosInstance, AxiosResponse} from "axios";
+
+import { getEnvironmentVariable } from "@langchain/core/utils/env";
+import { Embeddings, EmbeddingsParams } from "@langchain/core/embeddings";
+import { chunkArray } from "@langchain/core/utils/chunk_array";
+
+/**
+ * The default model name to use for generating embeddings.
+ */
+const DEFAULT_MODEL_NAME = "sentence-transformers/clip-ViT-B-32";
+
+/**
+ * The default batch size to use for generating embeddings.
+ * This is limited by the DeepInfra API to a maximum of 1024.
+ */
+const DEFAULT_BATCH_SIZE = 1024;
+
+/**
+ * Environment variable name for the DeepInfra API token.
+ */
+const API_TOKEN_ENV_VAR = "DEEPINFRA_API_TOKEN";
+
+
+export interface DeepInfraEmbeddingsRequest {
+  inputs: string[];
+  normalize?: boolean;
+  image?: string;
+  webhook?: string;
+}
+
+
+/**
+ * Input parameters for the DeepInfra embeddings
+ */
+export interface DeepInfraEmbeddingsParams extends EmbeddingsParams {
+
+  /**
+   * The API token to use for authentication.
+   * If not provided, it will be read from the `DEEPINFRA_API_TOKEN` environment variable.
+   */
+  apiToken?: string;
+
+  /**
+   * The model ID to use for generating completions.
+   * Default: `sentence-transformers/clip-ViT-B-32`
+   */
+  modelName?: string;
+
+  /**
+   * The maximum number of texts to embed in a single request. This is
+   * limited by the DeepInfra API to a maximum of 1024.
+   */
+  batchSize?: number;
+}
+
+/**
+ * Response from the DeepInfra embeddings API.
+ */
+export interface DeepInfraEmbeddingsResponse {
+  /**
+   * The embeddings generated for the input texts.
+   */
+  embeddings: number[][];
+  /**
+   * The number of tokens in the input texts.
+   */
+  input_tokens: number;
+  /**
+   * The status of the inference.
+   */
+  request_id?: string;
+}
+
+
+/**
+ * A class for generating embeddings using the Cohere API.
+ * @example
+ * ```typescript
+ * // Embed a query using the CohereEmbeddings class
+ * const model = new ChatOpenAI();
+ * const res = await model.embedQuery(
+ *   "What would be a good company name for a company that makes colorful socks?",
+ * );
+ * console.log({ res });
+ * ```
+ */
+export class DeepInfraEmbeddings
+  extends Embeddings
+  implements DeepInfraEmbeddingsParams
+{
+
+  private client: AxiosInstance;
+
+  apiToken: string;
+
+  batchSize: number;
+
+  modelName: string;
+
+
+  /**
+   * Constructor for the CohereEmbeddings class.
+   * @param fields - An optional object with properties to configure the instance.
+   */
+  constructor(
+    fields?: Partial<DeepInfraEmbeddingsParams> & {
+      verbose?: boolean;
+    }
+  ) {
+    const fieldsWithDefaults = {
+      modelName: DEFAULT_MODEL_NAME,
+      batchSize: DEFAULT_BATCH_SIZE,
+      ...fields };
+
+    super(fieldsWithDefaults);
+
+    const apiKey =
+      fieldsWithDefaults?.apiToken || getEnvironmentVariable(API_TOKEN_ENV_VAR);
+
+    if (!apiKey) {
+      throw new Error("DeepInfra API token not found");
+    }
+
+    this.modelName = fieldsWithDefaults?.modelName ?? this.modelName;
+    this.batchSize = fieldsWithDefaults?.batchSize ?? this.batchSize;
+    this.apiToken = apiKey;
+
+  }
+
+  /**
+   * Generates embeddings for an array of texts.
+   * @param inputs - An array of strings to generate embeddings for.
+   * @returns A Promise that resolves to an array of embeddings.
+   */
+  async embedDocuments(inputs: string[]): Promise<number[][]> {
+    await this.maybeInitClient();
+
+    const batches = chunkArray(inputs, this.batchSize) as string[][];
+
+    const batchRequests = batches.map((batch : string[]) =>
+      this.embeddingWithRetry({
+        inputs: batch,
+      })
+    );
+
+    const batchResponses = await Promise.all(batchRequests);
+
+    const out: number[][] = [];
+
+    for (let i = 0; i < batchResponses.length; i += 1) {
+      const batch = batches[i];
+      const { embeddings } = batchResponses[i];
+      for (let j = 0; j < batch.length; j += 1) {
+        out.push(embeddings[j]);
+      }
+    }
+
+    return out;
+  }
+
+  /**
+   * Generates an embedding for a single text.
+   * @param text - A string to generate an embedding for.
+   * @returns A Promise that resolves to an array of numbers representing the embedding.
+   */
+  async embedQuery(text: string): Promise<number[]> {
+    await this.maybeInitClient();
+
+    const {embeddings} = await this.embeddingWithRetry({
+      inputs: [text],
+    });
+    return embeddings[0];
+  }
+
+  /**
+   * Generates embeddings with retry capabilities.
+   * @param request - An object containing the request parameters for generating embeddings.
+   * @returns A Promise that resolves to the API response.
+   */
+  private async embeddingWithRetry(
+    request: DeepInfraEmbeddingsRequest
+  ): Promise<DeepInfraEmbeddingsResponse> {
+    this.maybeInitClient();
+    const response = await this.caller.call(this.client.post.bind(this.client,""), request);
+    return (response as AxiosResponse<DeepInfraEmbeddingsResponse>).data;
+  }
+
+  /**
+   * Initializes the DeepInfra client if it hasn't been initialized already.
+   */
+  private maybeInitClient() {
+    if (!this.client) {
+
+      this.client = axios.default.create({
+        baseURL: `https://api.deepinfra.com/v1/inference/${this.modelName}`,
+        headers: {
+          Authorization: `Bearer ${this.apiToken}`,
+          ContentType: "application/json",
+        },
+      });
+    }
+  }
+
+  /** @ignore */
+  static async imports(): Promise<{}> {
+    // Axios has already been defined as dependency in the package.json
+    // so we can use it here without importing it.
+    return {};
+  }
+}

libs/langchain-community/src/embeddings/tests/deepinfra.int.test.ts

@@ -0,0 +1,34 @@
+import { test, expect } from "@jest/globals";
+import { DeepInfraEmbeddings } from "../deepinfra.js";
+
+test("Test DeepInfraEmbeddings.embedQuery", async () => {
+  const embeddings = new DeepInfraEmbeddings();
+  const res = await embeddings.embedQuery("Hello world");
+  expect(typeof res[0]).toBe("number");
+});
+
+test("Test DeepInfraEmbeddings.embedDocuments", async () => {
+  const embeddings = new DeepInfraEmbeddings();
+  const res = await embeddings.embedDocuments(["Hello world", "Bye bye"]);
+  expect(res).toHaveLength(2);
+  expect(typeof res[0][0]).toBe("number");
+  expect(typeof res[1][0]).toBe("number");
+});
+
+test("Test DeepInfraEmbeddings concurrency", async () => {
+  const embeddings = new DeepInfraEmbeddings({
+    batchSize: 1,
+  });
+  const res = await embeddings.embedDocuments([
+    "Hello world",
+    "Bye bye",
+    "we need",
+    "at least",
+    "six documents",
+    "to test concurrency"
+  ]);
+  expect(res).toHaveLength(6);
+  expect(res.find((embedding) => typeof embedding[0] !== "number")).toBe(
+    undefined
+  );
+});