Update code.gs

code.gs

@@ -1,6 +1,9 @@
-const API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.3";
+const DEFAULT_SYSTEM_PROMPT = 'You are a helpful and honest assistant. Please, respond concisely and truthfully.';
+const DEFAULT_MODEL = "mistralai/Mistral-7B-Instruct-v0.3";
 
-// Function to display a menu in Google Sheets
+/**
+ * Adds a custom menu to Google Sheets to allow users to input their Hugging Face API key.
+ */
 function onOpen() {
   const ui = SpreadsheetApp.getUi();
   ui.createMenu('Hugging Sheets')
@@ -8,63 +11,94 @@ function onOpen() {
     .addToUi();
 }
 
-// Function to prompt the user for their Hugging Face API key
+/**
+ * Prompts the user to enter their Hugging Face API key and saves it in the script properties.
+ */
 function showApiKeyPrompt() {
   const ui = SpreadsheetApp.getUi();
   const response = ui.prompt('Enter your Hugging Face API Key:');
   if (response.getSelectedButton() == ui.Button.OK) {
-    const apiKey = response.getResponseText();
-    PropertiesService.getScriptProperties().setProperty('HF_API_KEY', apiKey);
-    ui.alert('API Key saved successfully!');
+    const apiKey = response.getResponseText().trim();
+    if (apiKey) {
+      PropertiesService.getScriptProperties().setProperty('HF_API_KEY', apiKey);
+      ui.alert('API Key saved successfully!');
+    } else {
+      ui.alert('API Key cannot be empty.');
+    }
   }
 }
 
-// Function to call the Hugging Face API
-function queryHuggingFace(prompt, model) {
+/**
+ * Sends a request to the Hugging Face API with the specified prompt and model.
+ * 
+ * @param {string} prompt - The input prompt to be sent to the model.
+ * @param {string} [model] - The model ID to query (e.g., 'mistralai/Mistral-7B-Instruct-v0.3').
+ * @param {string} [systemPrompt=DEFAULT_SYSTEM_PROMPT] - The system prompt to customize the assistant's behavior.
+ * @returns {object} JSON response from the Hugging Face API.
+ * @throws {Error} If the API key is not set or if the API request fails.
+ */
+function queryHuggingFace(prompt, model = DEFAULT_MODEL, systemPrompt = DEFAULT_SYSTEM_PROMPT) {
   const apiKey = PropertiesService.getScriptProperties().getProperty('HF_API_KEY');
   if (!apiKey) {
     throw new Error('Please enter your Hugging Face API key using the menu.');
   }
 
-  const url = `https://api-inference.huggingface.co/models/${model}`;
-  const headers = {
-    "Authorization": `Bearer ${apiKey}`,
-    "Content-Type": "application/json"
-  };
-
-
-  const formattedPrompt = `<s> [INST] You are a helpful and honest assistant. Please, respond concisely and truthfully. [/INST] ${prompt} </s>`;
+  const formattedPrompt = `<s> [INST] ${systemPrompt} [/INST] ${prompt} </s>`;
   const payload = {
     "inputs": formattedPrompt
   };
 
+  const payloadString = JSON.stringify(payload);
 
+  const url = `https://api-inference.huggingface.co/models/${model}`;
   const options = {
-    "method": "post",
-    "headers": headers,
-    "payload": JSON.stringify(payload)
+    'method': 'post',
+    'contentType': 'application/json',
+    'headers': {
+      'Authorization': `Bearer ${apiKey}`
+    },
+    'payload': payloadString
   };
 
-  const response = UrlFetchApp.fetch(url, options);
-  const json = JSON.parse(response.getContentText());
-  
-  return json;
+  try {
+    const response = UrlFetchApp.fetch(url, options);
+    const json = JSON.parse(response.getContentText());
+    return json;
+  } catch (error) {
+    throw new Error(`Failed to fetch data from Hugging Face API: ${error.message}`);
+  }
 }
 
-// Function to create the custom formula "=HF(prompt, model)"
-function HF(prompt, model) {
+/**
+ * Custom Google Sheets formula to query the Hugging Face API and return the generated text.
+ * Function to create the custom formula =HF(prompt, model, [systemPrompt])
+ * 
+ * @param {string} prompt - The input prompt to be sent to the model.
+ * @param {string} [model] - The model ID to query (e.g., 'mistralai/Mistral-7B-Instruct-v0.3').
+ * @param {string} [systemPrompt] - The system prompt to customize the assistant's behavior. Defaults to DEFAULT_SYSTEM_PROMPT.
+ * @returns {string} The generated output text from the Hugging Face API, or an error message if the request fails.
+ */
+function HF(prompt, model = DEFAULT_MODEL, systemPrompt = DEFAULT_SYSTEM_PROMPT) {
   try {
-    const response = queryHuggingFace(prompt, model);
-    const fullResponse = response[0].generated_text; // Adjust based on the actual response structure
-    // Extract the part of the response after the prompt
-    const generatedOutput = fullResponse.split(`</s>`).pop().trim();
-    return generatedOutput;
+    const response = queryHuggingFace(prompt, model, systemPrompt);
+    if (response && response.length > 0 && response[0].generated_text) {
+      const fullResponse = response[0].generated_text;
+      // Extract the part of the response after the prompt
+      const generatedOutput = fullResponse.split(`</s>`).pop().trim();
+      return generatedOutput;
+    } else {
+      return 'Error: Invalid response structure from Hugging Face API.';
+    }
   } catch (error) {
     return `Error: ${error.message}`;
   }
 }
 
-// Add the formula to Google Sheets
+/**
+ * Add the formula to Google Sheets.
+ * 
+ * @param {object} e - The event object.
+ */
 function onInstall(e) {
   onOpen(e);
   const formula = SpreadsheetApp.newUserDefinedFunctionBuilder()
@@ -72,4 +106,4 @@ function onInstall(e) {
     .setFunction('HF')
     .build();
   SpreadsheetApp.installUserDefinedFunction(formula);
-}
+}