Add Vercel SDK to node package to support multiple models

node-zerox/src/index.ts

@@ -5,7 +5,7 @@ import {
   formatMarkdown,
   isString,
 } from "./utils";
-import { getCompletion } from "./openAI";
+import { getCompletion } from "./processor";
 import { ModelOptions, ZeroxArgs, ZeroxOutput } from "./types";
 import { validateLLMParams } from "./utils";
 import fs from "fs-extra";
@@ -14,13 +14,13 @@ import path from "path";
 import pLimit, { Limit } from "p-limit";
 
 export const zerox = async ({
+  apiKey = "",
   cleanup = true,
   concurrency = 10,
   filePath,
   llmParams = {},
   maintainFormat = false,
   model = ModelOptions.gpt_4o_mini,
-  openaiAPIKey = "",
   outputDir,
   pagesToConvertAsImages = -1,
   tempDir = os.tmpdir(),
@@ -34,9 +34,6 @@ export const zerox = async ({
   llmParams = validateLLMParams(llmParams);
 
   // Validators
-  if (!openaiAPIKey || !openaiAPIKey.length) {
-    throw new Error("Missing OpenAI API key");
-  }
   if (!filePath || !filePath.length) {
     throw new Error("Missing file path");
   }
@@ -96,7 +93,7 @@ export const zerox = async ({
       const imagePath = path.join(tempDirectory, image);
       try {
         const { content, inputTokens, outputTokens } = await getCompletion({
-          apiKey: openaiAPIKey,
+          apiKey,
           imagePath,
           llmParams,
           maintainFormat,
@@ -122,7 +119,7 @@ export const zerox = async ({
       const imagePath = path.join(tempDirectory, image);
       try {
         const { content, inputTokens, outputTokens } = await getCompletion({
-          apiKey: openaiAPIKey,
+          apiKey,
           imagePath,
           llmParams,
           maintainFormat,

node-zerox/src/models.ts

@@ -0,0 +1,53 @@
+import { createAmazonBedrock } from "@ai-sdk/amazon-bedrock";
+import { createAnthropic } from "@ai-sdk/anthropic";
+import { createGoogleGenerativeAI } from "@ai-sdk/google";
+import { createMistral } from "@ai-sdk/mistral";
+import { createOpenAI } from "@ai-sdk/openai";
+
+const MODEL_PROVIDERS = {
+  anthropic: {
+    models: [
+      "claude-3-5-sonnet-20240620",
+      "claude-3-haiku-20240307",
+      "claude-3-opus-20240229",
+      "claude-3-sonnet-20240229",
+    ],
+    provider: createAnthropic,
+  },
+  bedrock: {
+    models: [
+      "anthropic.claude-3-5-sonnet-20240620-v1:0",
+      "anthropic.claude-3-haiku-20240307-v1:0",
+      "anthropic.claude-3-opus-20240229-v1:0",
+      "anthropic.claude-3-sonnet-20240229-v1:0",
+    ],
+    provider: createAmazonBedrock,
+  },
+  gemini: {
+    models: [
+      "gemini-1.5-flash-latest",
+      "gemini-1.5-flash",
+      "gemini-1.5-pro-latest",
+      "gemini-1.5-pro",
+    ],
+    provider: createGoogleGenerativeAI,
+  },
+  gpt: {
+    models: ["gpt-4-turbo", "gpt-4o-mini", "gpt-4o"],
+    provider: createOpenAI,
+  },
+  mistral: {
+    models: ["pixtral-12b-2409"],
+    provider: createMistral,
+  },
+};
+
+export const createProviderInstance = (model: string, apiKey: string) => {
+  const foundProvider = Object.values(MODEL_PROVIDERS).find((group) =>
+    group.models.includes(model)
+  );
+  if (foundProvider) {
+    return foundProvider.provider({ apiKey });
+  }
+  throw new Error(`Model '${model}' does not support image inputs`);
+};

node-zerox/src/openAI.ts → node-zerox/src/processor.ts

@@ -1,6 +1,7 @@
 import { CompletionArgs, CompletionResponse } from "./types";
 import { convertKeysToSnakeCase, encodeImageToBase64 } from "./utils";
-import axios from "axios";
+import { createProviderInstance } from "./models";
+import { generateText } from "ai";
 
 export const getCompletion = async ({
   apiKey,
@@ -41,27 +42,16 @@ export const getCompletion = async ({
   });
 
   try {
-    const response = await axios.post(
-      "https://api.openai.com/v1/chat/completions",
-      {
-        messages,
-        model,
-        ...convertKeysToSnakeCase(llmParams ?? null),
-      },
-      {
-        headers: {
-          Authorization: `Bearer ${apiKey}`,
-          "Content-Type": "application/json",
-        },
-      }
-    );
-
-    const data = response.data;
-
+    const providerInstance = createProviderInstance(model, apiKey);
+    const { text, usage } = await generateText({
+      model: providerInstance(model),
+      messages,
+      ...convertKeysToSnakeCase(llmParams ?? null),
+    });
     return {
-      content: data.choices[0].message.content,
-      inputTokens: data.usage.prompt_tokens,
-      outputTokens: data.usage.completion_tokens,
+      content: text,
+      inputTokens: usage.promptTokens,
+      outputTokens: usage.completionTokens,
     };
   } catch (err) {
     console.error("Error in OpenAI completion", err);

node-zerox/src/types.ts

@@ -1,11 +1,11 @@
 export interface ZeroxArgs {
+  apiKey?: string;
   cleanup?: boolean;
   concurrency?: number;
   filePath: string;
   llmParams?: LLMParams;
   maintainFormat?: boolean;
   model?: ModelOptions;
-  openaiAPIKey?: string;
   outputDir?: string;
   pagesToConvertAsImages?: number | number[];
   tempDir?: string;