Refactor; add new gemini model

htadashi · Aug 3, 2024 · 1483887 · 1483887
1 parent e41d01d
commit 1483887
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 26 deletions.
diff --git a/modules/gemini.js b/modules/gemini.js
@@ -1,4 +1,4 @@
-export function buildRequestGemini(request_params, apiKey) {
+export function buildRequestGemini(request_params, apiKey, model) {
 
     const prompt = request_params.prompt;
     const contents = {
@@ -31,7 +31,7 @@ export function buildRequestGemini(request_params, apiKey) {
         tool_config: tool_config
     });
     const request = {
-        endpoint: `https://generativelanguage.googleapis.com/v1beta/models/gemini-pro:generateContent?key=${apiKey}`,
+        endpoint: `https://generativelanguage.googleapis.com/v1beta/models/${model}:generateContent?key=${apiKey}`,
         options: {
             method: "POST",
             headers: {
@@ -48,7 +48,7 @@ export function parseResponseGemini(data) {
 
     const parsedResponse = {
         function_used: data.candidates[0].content.parts[0].functionCall.name,
-        event: JSON.parse(data.candidates[0].content.parts[0].functionCall.args)
+        event: data.candidates[0].content.parts[0].functionCall.args
     }
 
     return parsedResponse;

diff --git a/modules/openai.js b/modules/openai.js
@@ -1,7 +1,5 @@
 export function buildRequestOpenAI(request_params, apiKey, model) {
 
-    console.log(request_params);
-
     const prompt = request_params.prompt;
     let tools = [];
     for (const func of request_params.functions) {
@@ -11,8 +9,6 @@ export function buildRequestOpenAI(request_params, apiKey, model) {
         });
     }
 
-    console.log(request_params.functions.length);
-
     let tool_choice;
     if (request_params.functions.length === 1) {
         tool_choice = { "type": "function", "function": { "name": request_params.functions[0].name } };

diff --git a/options.html b/options.html
@@ -47,9 +47,10 @@
                     <select id="model-select">
                         <option value="gpt-3.5-turbo">GPT-3.5 Turbo</option>
                         <option value="gpt-4o">GPT-4o</option>
-                        <option value="gemini">Gemini 1.0</option>
+                        <option value="gemini-pro">Gemini 1.0 Pro</option>
+                        <option value="gemini-1.5-flash-latest">Gemini 1.5 Flash</option>
                     </select>
-                    <label>Choose default AI Model</label>
+                    <label>AI Model to be used:</label>
                 </div>
             </div>
             <div class="row">

diff --git a/service_worker.js b/service_worker.js
@@ -4,6 +4,30 @@ import { GCalLink, iCalDownload, autoSelect } from "./modules/prompts.js";
 
 import { isAllDayEvent } from "./modules/util.js";
 
+
+const modelHandlers = {
+    "gpt-3.5-turbo": { build: buildRequestOpenAI, parse: parseResponseOpenAI },
+    "gpt-4o": { build: buildRequestOpenAI, parse: parseResponseOpenAI },
+    "gemini-pro": { build: buildRequestGemini, parse: parseResponseGemini },
+    "gemini-1.5-flash-latest": { build: buildRequestGemini, parse: parseResponseGemini },
+};
+
+function buildRequest(request_params, apiKey, model) {
+    const handler = modelHandlers[model];
+    if (handler) {
+        return handler.build(request_params, apiKey, model);
+    }
+    throw new Error(`Unsupported model: ${model}`);
+}
+
+function parseResponse(data, model) {
+    const handler = modelHandlers[model];
+    if (handler) {
+        return handler.parse(data);
+    }
+    throw new Error(`Unsupported model: ${model}`);
+}
+
 chrome.runtime.onInstalled.addListener(() => {
     chrome.contextMenus.create({
         id: "create-gcal-url",
@@ -45,15 +69,7 @@ chrome.contextMenus.onClicked.addListener((info, tab) => {
                     request_params = autoSelect(selectedText);
                 }
 
-                console.log(request_params);
-
-                let request;
-                if (model === "gpt-3.5-turbo" || model === "gpt-4o") {
-                    request = buildRequestOpenAI(request_params, apiKey, model);
-                } else if (model === "gemini") {
-                    request = buildRequestGemini(request_params, apiKey);
-                }
-
+                const request = buildRequest(request_params, apiKey, model);
                 fetch(request.endpoint, request.options)
                     .then(response => response.json())
                     .then(data => {
@@ -64,16 +80,9 @@ chrome.contextMenus.onClicked.addListener((info, tab) => {
                             throw new Error("Invalid API key");
                         }
 
-                        let parsedResponse;
-                        if (model === "gpt-3.5-turbo" || model === "gpt-4o") {
-                            parsedResponse = parseResponseOpenAI(data);
-                        } else if (model === "gemini") {
-                            parsedResponse = parseResponseGemini(data);
-                        }
+                        const parsedResponse = parseResponse(data, model);
                         const { function_used, event } = parsedResponse;
 
-                        console.log(event);
-
                         if (function_used === "get_event_information") {
                             // Format the dates
                             const startDate = event.start_date;