Replaced ai-sdk/google

core/dist/types.d.ts

@@ -3,12 +3,14 @@ export declare enum OpenAIModels {
     GPT_4O_MINI = "gpt-4o-mini"
 }
 export declare enum LocalModels {
-    LLAVA = "llava",
     LLAMA3_2_VISION = "llama3.2-vision"
 }
 export declare enum GoogleModels {
     GEMINI_2_FLASH = "gemini-2.0-flash-001",
-    GEMINI_2_FLASH_LITE = "gemini-2.0-flash-lite-preview-02-05"
+    GEMINI_2_FLASH_LITE = "gemini-2.0-flash-lite-preview-02-05",
+    GEMINI_1_5_FLASH = "gemini-1.5-flash",
+    GEMINI_1_5_FLASH_8B = "gemini-1.5-flash-8b",
+    GEMINI_1_5_PRO = "gemini-1.5-pro"
 }
 export type ModelOptions = OpenAIModels | GoogleModels | LocalModels;
 export interface DocumindArgs {

core/dist/types.js

@@ -8,11 +8,14 @@ var OpenAIModels;
 })(OpenAIModels || (exports.OpenAIModels = OpenAIModels = {}));
 var LocalModels;
 (function (LocalModels) {
-    LocalModels["LLAVA"] = "llava";
+    //LLAVA = "llava",
     LocalModels["LLAMA3_2_VISION"] = "llama3.2-vision";
 })(LocalModels || (exports.LocalModels = LocalModels = {}));
 var GoogleModels;
 (function (GoogleModels) {
     GoogleModels["GEMINI_2_FLASH"] = "gemini-2.0-flash-001";
     GoogleModels["GEMINI_2_FLASH_LITE"] = "gemini-2.0-flash-lite-preview-02-05";
+    GoogleModels["GEMINI_1_5_FLASH"] = "gemini-1.5-flash";
+    GoogleModels["GEMINI_1_5_FLASH_8B"] = "gemini-1.5-flash-8b";
+    GoogleModels["GEMINI_1_5_PRO"] = "gemini-1.5-pro";
 })(GoogleModels || (exports.GoogleModels = GoogleModels = {}));

core/src/types.ts

@@ -4,13 +4,16 @@ export enum OpenAIModels {
 }
 
 export enum LocalModels {
-  LLAVA = "llava",
+  //LLAVA = "llava",
   LLAMA3_2_VISION = "llama3.2-vision",
 }
 
 export enum GoogleModels {
   GEMINI_2_FLASH = "gemini-2.0-flash-001",
   GEMINI_2_FLASH_LITE = "gemini-2.0-flash-lite-preview-02-05",
+  GEMINI_1_5_FLASH = "gemini-1.5-flash",
+  GEMINI_1_5_FLASH_8B = "gemini-1.5-flash-8b",
+  GEMINI_1_5_PRO = "gemini-1.5-pro",
 }
 
 export type ModelOptions = OpenAIModels | GoogleModels | LocalModels;

extractor/package.json

@@ -9,14 +9,16 @@
   },
   "dependencies": {
     "@ai-sdk/google": "^1.1.14",
+    "@google/generative-ai": "^0.22.0",
     "ai": "^4.1.41",
     "axios": "^1.7.7",
     "core": "*",
     "dotenv": "^16.4.5",
     "openai": "^4.68.4",
     "pdf-lib": "^1.17.1",
     "uuid": "^11.0.2",
-    "zod": "^3.23.8"
+    "zod": "^3.23.8",
+    "zod-to-json-schema": "^3.24.2"
   },
   "devDependencies": {
     "nodemon": "^3.1.7"

extractor/src/autoschema/autogenerateSchema.js

@@ -42,7 +42,8 @@ async function blanketSchema(markdown, model) {
   const result = await extraction({
     markdown,
     zodSchema: schemaToUse,
-    prompt: AUTO_SCHEMA_PROMPT(markdown), 
+    prompt: AUTO_SCHEMA_PROMPT(markdown),
+    model: model, 
   });
 
   if (!result || !result.fields) {
@@ -67,7 +68,8 @@ async function instructionBasedSchema(markdown, model, instructions) {
   const instructionFields = await extraction({
     markdown: instructions,
     zodSchema: instructionsZod,
-    prompt: instructionPrompt
+    prompt: instructionPrompt,
+    model: model,
   });
 
   if (!instructionFields || !instructionFields.fields) {
@@ -81,7 +83,8 @@ async function instructionBasedSchema(markdown, model, instructions) {
   const result = await extraction({
     markdown,
     zodSchema: schemaToUse,
-    prompt: INSTRUCTIONS_SCHEMA_PROMPT(markdown, data), 
+    prompt: INSTRUCTIONS_SCHEMA_PROMPT(markdown, data),
+    model: model,
   });
 
   if (!result || !result.fields) {

extractor/src/extractors/google.js

@@ -1,25 +1,49 @@
-import { createGoogleGenerativeAI } from '@ai-sdk/google';
-import { generateObject } from 'ai';
+import { GoogleGenerativeAI } from "@google/generative-ai";
+import { zodToJsonSchema } from "zod-to-json-schema";
 
-export const googleExtractor = async ({ markdown, zodSchema, prompt }) => {
+const genAI = new GoogleGenerativeAI(process.env.GEMINI_API_KEY);
+
+export const googleExtractor = async ({ markdown, zodSchema, prompt, model }) => {
     if (!process.env.GEMINI_API_KEY) {
       throw new Error("Missing GEMINI_API_KEY");
     }
 
-const google = createGoogleGenerativeAI({
-    apiKey: process.env.GEMINI_API_KEY,
-});
-const googleModel = "gemini-2.0-flash-001"
-
-const completion = await generateObject({
-    model: google(googleModel, {
-      structuredOutputs: false,
-    }),
-    schema: zodSchema,
-    prompt: markdown,
-    system: prompt,
-  });
-
-  const event = completion.object;
-  return event;
-}
+const googleModel = model
+
+// Convert Zod schema to JSON schema
+let jsonSchema = zodToJsonSchema(zodSchema);
+
+// Remove additionalProperties and $schema keys
+const removeKeys = (obj) => {
+    if (Array.isArray(obj)) {
+        return obj.map(removeKeys);
+    } else if (typeof obj === "object" && obj !== null) {
+        return Object.fromEntries(
+            Object.entries(obj)
+                .filter(([key]) => key !== "additionalProperties" && key !== "$schema")
+                .map(([key, value]) => [key, removeKeys(value)])
+        );
+    }
+    return obj;
+};
+
+jsonSchema = removeKeys(jsonSchema);
+
+const modelToUse = genAI.getGenerativeModel({
+    model: googleModel,
+    systemInstruction: prompt,
+    generationConfig: {
+      responseMimeType: "application/json",
+      responseSchema: jsonSchema,
+      },
+    });
+
+const result = await modelToUse.generateContent(
+    markdown,
+  );
+
+//console.log(result.response.text());
+const event = result.response.text();
+return event;
+}
+

extractor/src/extractors/index.js

@@ -3,8 +3,14 @@ import { openAIExtractor } from "./openAI.js";
 import { googleExtractor } from "./google.js";
 
 export const OpenAIModels = ["gpt-4o", "gpt-4o-mini"];
-export const LocalModels = ["llava", "llama3.2-vision"];
-export const GoogleModels = ["gemini-2.0-flash-001", "gemini-2.0-flash-lite-preview-02-05"]
+export const LocalModels = ["llama3.2-vision"];
+export const GoogleModels = [
+  "gemini-2.0-flash-001", 
+  "gemini-2.0-flash-lite-preview-02-05", 
+  "gemini-1.5-flash",
+  "gemini-1.5-flash-8b",
+ "gemini-1.5-pro"
+];
 
 export function getExtractor(model) {
   if (OpenAIModels.includes(model)) {

extractor/src/extractors/ollama.js

@@ -1,7 +1,7 @@
 import OpenAI from "openai";
 import { zodResponseFormat } from "openai/helpers/zod";
 
-export const ollamaExtractor = async ({ markdown, zodSchema, prompt }) => {
+export const ollamaExtractor = async ({ markdown, zodSchema, prompt, model }) => {
   if (!process.env.BASE_URL) {
     throw new Error("Missing BASE_URL");
   }
@@ -11,7 +11,7 @@ export const ollamaExtractor = async ({ markdown, zodSchema, prompt }) => {
     apiKey: 'ollama'
    });
 
-  const ollamaModel = "llama3.1";
+  const ollamaModel = model;
 
   const completion = await openai.beta.chat.completions.parse({
     model: ollamaModel,

extractor/src/extractors/openAI.js

@@ -1,13 +1,13 @@
 import OpenAI from "openai";
 import { zodResponseFormat } from "openai/helpers/zod";
 
-export const openAIExtractor = async ({ markdown, zodSchema, prompt }) => {
+export const openAIExtractor = async ({ markdown, zodSchema, prompt, model }) => {
   if (!process.env.OPENAI_API_KEY) {
     throw new Error("Missing OPENAI_API_KEY");
   }
 
   const openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY });
-  const openAiModel = "gpt-4o-2024-08-06";
+  const openAiModel = model;
 
   const completion = await openai.beta.chat.completions.parse({
     model: openAiModel,

extractor/src/services/extract.js

@@ -59,6 +59,7 @@ export async function extract({ file, schema, template, model, autoSchema }) {
       markdown,
       zodSchema: dynamicZodSchema,
       prompt: BASE_EXTRACTION_PROMPT,
+      model: defaultModel,
     });
 
     return {