Merge pull request #195 from Portkey-AI/feat/nomic-embeddings

feat: add nomic embeddings integration
Portkey-AI · Feb 13, 2024 · 82ab583 · 82ab583
2 parents 90bfef6 + 9b796f2
commit 82ab583
Show file tree

Hide file tree

Showing 7 changed files with 149 additions and 3 deletions.
diff --git a/src/globals.ts b/src/globals.ts
@@ -31,6 +31,7 @@ export const GOOGLE: string = "google";
 export const PERPLEXITY_AI: string = "perplexity-ai";
 export const MISTRAL_AI: string = "mistral-ai";
 export const DEEPINFRA: string = "deepinfra";
+export const NOMIC: string = "nomic";
 
 export const providersWithStreamingSupport = [OPEN_AI, AZURE_OPEN_AI, ANTHROPIC, COHERE];
 export const allowedProxyProviders = [OPEN_AI, COHERE, AZURE_OPEN_AI, ANTHROPIC];

diff --git a/src/middlewares/requestValidator/index.ts b/src/middlewares/requestValidator/index.ts
@@ -13,6 +13,7 @@ import {
     POWERED_BY,
     TOGETHER_AI,
     DEEPINFRA,
+    NOMIC,
 } from "../../globals";
 import { configSchema } from "./schema/config";
 
@@ -61,7 +62,7 @@ export const requestValidator = (c: Context, next: any) => {
     }
     if (
         requestHeaders[`x-${POWERED_BY}-provider`] &&
-        ![OPEN_AI, AZURE_OPEN_AI, COHERE, ANTHROPIC, ANYSCALE, PALM, TOGETHER_AI, GOOGLE, MISTRAL_AI, PERPLEXITY_AI, DEEPINFRA].includes(
+        ![OPEN_AI, AZURE_OPEN_AI, COHERE, ANTHROPIC, ANYSCALE, PALM, TOGETHER_AI, GOOGLE, MISTRAL_AI, PERPLEXITY_AI, DEEPINFRA, NOMIC].includes(
             requestHeaders[`x-${POWERED_BY}-provider`]
         )
     ) {

diff --git a/src/middlewares/requestValidator/schema/config.ts b/src/middlewares/requestValidator/schema/config.ts
@@ -10,6 +10,7 @@ import {
     PERPLEXITY_AI,
     TOGETHER_AI,
     DEEPINFRA,
+    NOMIC,
 } from "../../../globals";
 
 export const configSchema: any = z
@@ -45,7 +46,8 @@ export const configSchema: any = z
                         GOOGLE,
                         PERPLEXITY_AI,
                         MISTRAL_AI,
-                        DEEPINFRA
+                        DEEPINFRA,
+                        NOMIC
                     ].includes(value),
                 {
                     message:

diff --git a/src/providers/index.ts b/src/providers/index.ts
@@ -5,6 +5,7 @@ import CohereConfig from "./cohere";
 import DeepInfraConfig from "./deepinfra";
 import GoogleConfig from "./google";
 import MistralAIConfig from "./mistral-ai";
+import NomicConfig from "./nomic";
 import OpenAIConfig from "./openai";
 import PalmAIConfig from "./palm";
 import PerplexityAIConfig from "./perplexity-ai";
@@ -22,7 +23,8 @@ const Providers: { [key: string]: ProviderConfigs } = {
   google: GoogleConfig,
   'perplexity-ai': PerplexityAIConfig,
   'mistral-ai': MistralAIConfig,
-  'deepinfra': DeepInfraConfig
+  'deepinfra': DeepInfraConfig,
+  nomic: NomicConfig
 };
 
 export default Providers;
diff --git a/src/providers/nomic/api.ts b/src/providers/nomic/api.ts
@@ -0,0 +1,11 @@
+import { ProviderAPIConfig } from "../types";
+
+const NomicAPIConfig: ProviderAPIConfig = {
+    baseURL: "https://api-atlas.nomic.ai/v1",
+    headers: (API_KEY: string) => {
+        return { Authorization: `Bearer ${API_KEY}` };
+    },
+    embed: "/embedding/text",
+};
+
+export default NomicAPIConfig;
diff --git a/src/providers/nomic/embed.ts b/src/providers/nomic/embed.ts
@@ -0,0 +1,116 @@
+import { NOMIC } from "../../globals";
+import { EmbedParams, EmbedResponse } from "../../types/embedRequestBody";
+import { ErrorResponse, ProviderConfig } from "../types";
+
+export const NomicEmbedConfig: ProviderConfig = {
+    model: {
+        param: "model",
+        required: true,
+        default: "nomic-embed-text-v1",
+    },
+    input: {
+        param: "texts",
+        required: true,
+        transform: (params: EmbedParams) => {
+            if (Array.isArray(params.input)) {
+                return params.input;
+            }
+
+            return [params.input];
+        },
+    },
+    task_type: {
+        param: "task_type"
+    }
+};
+
+interface NomicEmbedResponse {
+    embeddings: number[][];
+    model: string;
+    usage: {
+        prompt_tokens: number;
+        total_tokens: number;
+    };
+}
+
+export interface NomicValidationErrorResponse {
+    detail: {
+        loc: Array<any>;
+        msg: string;
+        type: string;
+    }[];
+}
+
+export interface NomicErrorResponse {
+    detail: string;
+}
+
+export const NomicErrorResponseTransform: (response: NomicValidationErrorResponse | NomicErrorResponse) => ErrorResponse = (response) => {
+    let firstError: Record<string, any> | undefined;
+    let errorField: string | null = null;
+    let errorMessage: string | undefined;
+    let errorType: string | null = null;
+
+    if (Array.isArray(response.detail)) {
+      [firstError] = response.detail;
+      errorField = firstError?.loc?.join(".") ?? "";
+      errorMessage = firstError.msg;
+      errorType = firstError.type;
+    } else {
+      errorMessage = response.detail;
+    }
+
+    return {
+      error: {
+        message: `${errorField ? `${errorField}: ` : ""}${errorMessage}`,
+        type: errorType,
+        param: null,
+        code: null,
+      },
+      provider: NOMIC,
+    } as ErrorResponse;
+  }
+
+
+
+export const NomicEmbedResponseTransform: (
+    response: NomicEmbedResponse | NomicValidationErrorResponse | NomicErrorResponse,
+    responseStatus: number
+) => EmbedResponse | ErrorResponse = (response, responseStatus) => {
+    if (
+        "detail" in response &&
+        responseStatus !== 200 &&
+        response.detail.length
+      ) {
+        return NomicErrorResponseTransform(response);
+    }
+
+    if ("embeddings" in response) {
+        return {
+            object: "list",
+            data: response.embeddings.map((d, index) => ({
+                object: "embedding",
+                embedding: d,
+                index: index,
+            })),
+            model: response.model,
+            usage: {
+                prompt_tokens: response.usage.prompt_tokens,
+                total_tokens: response.usage.total_tokens,
+            },
+            provider: NOMIC,
+        };
+    }
+
+    return {
+        error: {
+            message: `Invalid response recieved from ${NOMIC}: ${JSON.stringify(
+                response
+            )}`,
+            type: null,
+            param: null,
+            code: null,
+        },
+        provider: NOMIC,
+    } as ErrorResponse;
+};
diff --git a/src/providers/nomic/index.ts b/src/providers/nomic/index.ts
@@ -0,0 +1,13 @@
+import { ProviderConfigs } from "../types";
+import NomicAPIConfig from "./api";
+import { NomicEmbedConfig, NomicEmbedResponseTransform } from "./embed";
+
+const NomicConfig: ProviderConfigs = {
+    embed: NomicEmbedConfig,
+    api: NomicAPIConfig,
+    responseTransforms: {
+        embed: NomicEmbedResponseTransform,
+    },
+};
+
+export default NomicConfig;