From b6b70b0fb2415f99b88c823ffcd26948b99a7c27 Mon Sep 17 00:00:00 2001
From: Alexander Alemayhu <alexander@alemayhu.com>
Date: Fri, 27 Dec 2024 05:03:50 +0100
Subject: [PATCH] refactor: reduce vertex duplication

Signed-off-by: Alexander Alemayhu <alexander@alemayhu.com>
---
 .../adapters/fileConversion/constants.ts      | 25 +++++++++++++++----
 .../fileConversion/convertImageToHTML.ts      | 19 ++------------
 .../fileConversion/convertPDFToHTML.ts        | 20 +++------------
 .../adapters/fileConversion/vertexAIUtils.ts  | 15 +++++++++++
 4 files changed, 40 insertions(+), 39 deletions(-)
 create mode 100644 src/infrastracture/adapters/fileConversion/vertexAIUtils.ts

diff --git a/src/infrastracture/adapters/fileConversion/constants.ts b/src/infrastracture/adapters/fileConversion/constants.ts
index 56f0972f..f312e06a 100644
--- a/src/infrastracture/adapters/fileConversion/constants.ts
+++ b/src/infrastracture/adapters/fileConversion/constants.ts
@@ -1,12 +1,16 @@
-import { HarmBlockThreshold, HarmCategory } from '@google-cloud/vertexai';
+import {
+  HarmBlockThreshold,
+  HarmCategory,
+  SafetySetting,
+} from '@google-cloud/vertexai';
 
-export const SAFETY_SETTINGS = [
+export const SAFETY_SETTINGS: SafetySetting[] = [
   {
-    category: HarmCategory.HARM_CATEGORY_HATE_SPEECH,
+    category: HarmCategory.HARM_CATEGORY_HARASSMENT,
     threshold: HarmBlockThreshold.BLOCK_NONE,
   },
   {
-    category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
+    category: HarmCategory.HARM_CATEGORY_HATE_SPEECH,
     threshold: HarmBlockThreshold.BLOCK_NONE,
   },
   {
@@ -14,7 +18,18 @@ export const SAFETY_SETTINGS = [
     threshold: HarmBlockThreshold.BLOCK_NONE,
   },
   {
-    category: HarmCategory.HARM_CATEGORY_HARASSMENT,
+    category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
     threshold: HarmBlockThreshold.BLOCK_NONE,
   },
 ];
+
+export const VERTEX_AI_CONFIG = {
+  project: 'notion-to-anki',
+  location: 'europe-west3',
+  model: 'gemini-1.5-pro-002',
+  generationConfig: {
+    maxOutputTokens: 8192,
+    temperature: 1,
+    topP: 0.95,
+  },
+};
diff --git a/src/infrastracture/adapters/fileConversion/convertImageToHTML.ts b/src/infrastracture/adapters/fileConversion/convertImageToHTML.ts
index c863d266..c4b802e7 100644
--- a/src/infrastracture/adapters/fileConversion/convertImageToHTML.ts
+++ b/src/infrastracture/adapters/fileConversion/convertImageToHTML.ts
@@ -1,5 +1,4 @@
-import { VertexAI } from '@google-cloud/vertexai';
-import { SAFETY_SETTINGS } from './constants';
+import { setupVertexAI } from './vertexAIUtils';
 
 /**
  * Google VertexAI is returning Markdown:
@@ -16,21 +15,7 @@ export function removeFirstAndLastLine(content: string): string {
 export const convertImageToHTML = async (
   imageData: string
 ): Promise<string> => {
-  const vertexAI = new VertexAI({
-    project: 'notion-to-anki',
-    location: 'europe-west3',
-  });
-  const model = 'gemini-1.5-pro-002';
-
-  const generativeModel = vertexAI.preview.getGenerativeModel({
-    model: model,
-    generationConfig: {
-      maxOutputTokens: 8192,
-      temperature: 1,
-      topP: 0.95,
-    },
-    safetySettings: SAFETY_SETTINGS,
-  });
+  const generativeModel = setupVertexAI();
 
   const text1 = {
     text: `Convert the text in this image to the following format for (every question is their own ul):
diff --git a/src/infrastracture/adapters/fileConversion/convertPDFToHTML.ts b/src/infrastracture/adapters/fileConversion/convertPDFToHTML.ts
index 78f9b28d..0bfa0f63 100644
--- a/src/infrastracture/adapters/fileConversion/convertPDFToHTML.ts
+++ b/src/infrastracture/adapters/fileConversion/convertPDFToHTML.ts
@@ -1,24 +1,10 @@
 import path from 'path';
 import fs from 'fs';
-
-import { GenerateContentRequest, VertexAI } from '@google-cloud/vertexai';
-import { SAFETY_SETTINGS } from './constants';
+import { GenerateContentRequest } from '@google-cloud/vertexai';
+import { setupVertexAI } from './vertexAIUtils';
 
 export const convertPDFToHTML = async (pdf: string): Promise<string> => {
-  const vertexAI = new VertexAI({
-    project: 'notion-to-anki',
-    location: 'europe-west3',
-  });
-  const model = 'gemini-1.5-pro-002';
-  const generativeModel = vertexAI.preview.getGenerativeModel({
-    model: model,
-    generationConfig: {
-      maxOutputTokens: 8192,
-      temperature: 1,
-      topP: 0.95,
-    },
-    safetySettings: SAFETY_SETTINGS,
-  });
+  const generativeModel = setupVertexAI();
 
   const document1 = {
     inlineData: {
diff --git a/src/infrastracture/adapters/fileConversion/vertexAIUtils.ts b/src/infrastracture/adapters/fileConversion/vertexAIUtils.ts
new file mode 100644
index 00000000..4f6ff6da
--- /dev/null
+++ b/src/infrastracture/adapters/fileConversion/vertexAIUtils.ts
@@ -0,0 +1,15 @@
+import { GenerativeModel, VertexAI } from '@google-cloud/vertexai';
+import { SAFETY_SETTINGS, VERTEX_AI_CONFIG } from './constants';
+
+export function setupVertexAI(): GenerativeModel {
+  const vertexAI = new VertexAI({
+    project: VERTEX_AI_CONFIG.project,
+    location: VERTEX_AI_CONFIG.location,
+  });
+
+  return vertexAI.getGenerativeModel({
+    model: VERTEX_AI_CONFIG.model,
+    generationConfig: VERTEX_AI_CONFIG.generationConfig,
+    safetySettings: SAFETY_SETTINGS,
+  });
+}