const API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.3";

// Function to display a menu in Google Sheets
function onOpen() {
  const ui = SpreadsheetApp.getUi();
  ui.createMenu('Hugging Sheets')
    .addItem('Enter your HF API Key', 'showApiKeyPrompt')
    .addToUi();
}

// Function to prompt the user for their Hugging Face API key
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!');
  }
}

// Function to call the Hugging Face API
function queryHuggingFace(prompt, model) {
  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 payload = {
    "inputs": formattedPrompt
  };


  const options = {
    "method": "post",
    "headers": headers,
    "payload": JSON.stringify(payload)
  };

  const response = UrlFetchApp.fetch(url, options);
  const json = JSON.parse(response.getContentText());
  
  return json;
}

// Function to create the custom formula "=HF(prompt, model)"
function HF(prompt, model) {
  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;
  } catch (error) {
    return `Error: ${error.message}`;
  }
}

// Add the formula to Google Sheets
function onInstall(e) {
  onOpen(e);
  const formula = SpreadsheetApp.newUserDefinedFunctionBuilder()
    .setName('HF')
    .setFunction('HF')
    .build();
  SpreadsheetApp.installUserDefinedFunction(formula);
}