Modifying Bigcode Request Method

background.js

@@ -17,11 +17,17 @@ chrome.runtime.onMessage.addListener((request, sender, sendResponse) => {
       });
       return true; // Required to use sendResponse asynchronously
     }
-    else if(request.type === "getmodelType"){
+    else if(request.type === "getModelType"){
       chrome.storage.sync.get("modelType", (data) => {
         sendResponse({ modelType: data.modelType });
       });
       return true; // Required to use sendResponse asynchronously
     }
+    else if(request.type === "getHuggingfaceApiKey"){
+      chrome.storage.sync.get("huggingfaceApiKey", (data) => {
+        sendResponse({ huggingfaceApiKey : data.huggingfaceApiKey});
+      })
+      return true;
+    }
   });
 

content.js

@@ -21,28 +21,30 @@ async function getChecked() {
 }
 async function getmodelType() {
   return new Promise((resolve) => {
-    chrome.runtime.sendMessage({ type: "getmodelType" }, (response) => {
+    chrome.runtime.sendMessage({ type: "getModelType" }, (response) => {
       resolve(response.modelType);
     });
   });
 }
+async function getHuggingfaceApiKey() {
+  return new Promise((resolve) => {
+    chrome.runtime.sendMessage({ type: "getHuggingfaceApiKey" }, (response) => {
+      resolve(response.huggingfaceApiKey)
+    })
+  })
+}
 
-// Use a regular expression to match the content between triple backticks
-const codeBlockRegex = /```([\s\S]*)```/;
-const codeHalfBlockRegex = /```([\s\S]*)/;
 
 // Function to send request to OpenAI API
 async function sendToOpenAI(prompt) {
   const apiKey = await getOpenAIKey();
   const modelType = await getmodelType();
-  if (!apiKey) {
-    // 总是忘记填写。。所以等了半天总是以为网络错误，改成了alert
-    alert("OpenAI API key not set."); 
-    return;
-  }else if(!modelType) {
-    alert("modelType not set."); 
+
+  if (!apiKey || !modelType) {
+    alert("OpenAI API key or modelType not set."); 
     return;
   }
+
   const response = await fetch("https://api.openai.com/v1/completions", {
     method: "POST",
     headers: {
@@ -52,56 +54,57 @@ async function sendToOpenAI(prompt) {
     body: JSON.stringify({
       model: modelType,
       prompt,
-      temperature: 0.1,
+      temperature: 0,
       max_tokens: 40,
     }),
-    // 添加一个最大请求时间，目前没做相关于超时的处理，但如果超时了，就会在请求的单元格展示error
     timeout: 30000
   });
 
   const data = await response.json();
-  const suggestion = data.choices && data.choices[0] && data.choices[0].text;
-
+  if(!(data.choices && data.choices[0])){
+    return null
+  }
 
-  // don't know how many "```" exists, It is related to the token and also related to the model
-  let count = (suggestion.match(/```/g) || []).length;
-  let code = ""
-  switch(count){
-    case 1: 
-      var match = suggestion.match(codeHalfBlockRegex)
-      code = match && match[1] ? match[1].replace(/^\n\s*/, '').replace(/\n.*$/, '').replace(/\u200B/g, ''):""
-      break;
-    case 2:
-      var match = suggestion.match(codeBlockRegex)
-      code = match && match[1] ? match[1].replace(/^\n\s*/, '').replace(/\n.*$/, '').replace(/\u200B/g, ''):""
-      break;
-    default:code = suggestion;
+  let suggestion = data.choices && data.choices[0] && data.choices[0].text;
+  // Remove invisible characters
+  suggestion = suggestion.replace(/\u200B/g, '');
+
+  // This is an example of a possible return: "    print('Hello World!')\n\nhello_world()\n\n§ Output… : ['Hello World!\\n']\n\n \n§ Markdown\n\n### Exercise"
+  const outPutIndex = suggestion.indexOf("\n\n§ Output")
+  if(outPutIndex == -1){
+    return suggestion
+  }else{
+    return suggestion.substring(0,outPutIndex)
   }
 
-  return code
 }
 
 async function sendToOtherService(code) {
   const url = await getOtherServiceUrl();
+  const token = await getHuggingfaceApiKey();
+
   if (!url) {
-    // 总是忘记填写。。所以等了半天总是以为网络错误，改成了alert
     alert("otherServiceUrl not set.");
     return;
   }
-  const prompt = "<start_jupyter><jupyter_code>" + code.replace(/\u200B/g, '')
+
+  const prompt = code.replace(/\u200B/g, '')
   const response = await fetch(url, {
     method: "POST",
     headers: {
-      "Content-Type": "application/json"
+      "Content-Type": "application/json",
+      "Authorization": `Bearer ${token}`,
     },
     body: JSON.stringify({
       inputs: prompt,
+      stream: false,
       parameters: {
         return_full_text: false
       }
     })
   })
   const data = await response.json();
+
   if (/^\n/.test(data)) {
     data = data.replace(/^\n/, '');
   }
@@ -110,38 +113,22 @@ async function sendToOtherService(code) {
 }
 
 
-
 async function getCodeCompletion(code) {
   const checked = await getChecked();
-  // 如果没有选择则弹出框提示
+  // If the user has not selected openai or bigcode
   if (!checked) {
     alert("The request method is not selected.");
     return;
   }
+
   if (checked == "openaiApiKey") {
-    return await getOpenAiWrapper(code)
+    return await sendToOpenAI(code)
   } else if (checked == "otherService") {
-    return await getOtherServiceUrlWrapper(code)
+    return await sendToOtherService(code)
   }
 }
 
-async function getOpenAiWrapper(code) {
-  const prompt = `
-  I have the following code:
-
-  \`\`\`
-  ${code}
-  \`\`\`
-
-  Please provide the missing code to complete the task. (do not include code that is already present in the prompt)
-  `;
 
-  return await sendToOpenAI(prompt);
-}
-
-async function getOtherServiceUrlWrapper(code) {
-  return await sendToOtherService(code)
-}
 
 // Code to be filled in after request completion
 let codeToFill = "";
@@ -193,6 +180,7 @@ const getActiveCellPointerCode = (activeCell) => {
 
     // Which line
     const lineIndex = Math.round(parseFloat(style.getPropertyValue('top')) / 17)
+
     // Obtain element for all line
     const linesElement = activeCell.getElementsByClassName('CodeMirror-line')
     // code dom element length in active line
@@ -204,39 +192,80 @@ const getActiveCellPointerCode = (activeCell) => {
     }
 
     for (let i = 0; i < linesElement.length; i++) {
-      if(i <= lineIndex) {
+
+      if(i < lineIndex) {
         leftContext += linesElement[i].textContent + "\n"
-      }else {
-        rightContext += linesElement[i].textContent  + "\n"
+      }else if(i == lineIndex){
+        leftContext += linesElement[i].textContent
+      }else{
+        if(i == linesElement.length-1){
+          rightContext += linesElement[i].textContent
+        }else{
+          rightContext += linesElement[i].textContent  + "\n"
+        }
       }
+
     }
 
     return [leftContext, rightContext]
 }
 
 
-
-function getCellContentText(activeCell) {
+function getCellContentTextRequiredForOpenAI(activeCell) {
   const cellElements = Array.from(document.querySelectorAll('.cell'));
   const activeCellIndex = cellElements.findIndex(cell => cell.contains(activeCell));
   // Check if there are at least 3 cells before the active cell
+  let codeContent = "";
+
+  // LeftContext refers to the left side of the pointer, and vice versa, If both are null, it is determined that the pointer is not at the far right
+  const [leftContext, rightContext] = getActiveCellPointerCode(activeCell)
+
+  if(!leftContext){
+    return null
+  }
 
+  // Iterate through the last 3 cells before the active cell
+  const startIndex = activeCellIndex - 3 < 0 ? 0 : activeCellIndex - 3;
+  for (let i = startIndex; i <= activeCellIndex; i++) {
+    if(i == activeCellIndex){
+      codeContent += leftContext
+      break
+    }else{
+      const cellElement = cellElements[i];
+      if (cellElement.classList.contains('code_cell')) {
+        codeContent += extractTextFromCell(cellElement);
+      }
+    }
+    codeContent += "\n"
+  }
+
+  return codeContent;
+}
+
+
+function getCellContentTextRequiredForBigCode(activeCell) {
+  const cellElements = Array.from(document.querySelectorAll('.cell'));
+  const activeCellIndex = cellElements.findIndex(cell => cell.contains(activeCell));
+  // Check if there are at least 3 cells before the active cell
   let combinedContent = "<start_jupyter>";
 
-  // LeftContext refers to the left side of the pointer, and vice versa, If both are null, it is determined that the pointer is not at the far right
+  // in active cell, LeftContext refers to the left side of the pointer, and vice versa, If both are null, it is determined that the pointer is not at the far right
   const [leftContext, rightContext] = getActiveCellPointerCode(activeCell)
 
   if(!leftContext && !rightContext){
     return null
   }
 
+  TODO: "The following code needs to add 'leftContext' and 'rightContext'"
   // Iterate through the last 3 cells before the active cell
   const startIndex = activeCellIndex - 3 < 0 ? 0 : activeCellIndex - 3;
+
   for (let i = startIndex; i <= activeCellIndex; i++) {
     const cellElement = cellElements[i];
 
     if (cellElement.classList.contains('code_cell')) {
       const code = extractTextFromCell(cellElement);
+
       combinedContent += `<jupyter_code>${code}`;
       const outputElement = cellElement.querySelector('.output_subarea');
       if (outputElement) {
@@ -249,21 +278,34 @@ function getCellContentText(activeCell) {
       const text = extractTextFromCell(cellElement);
       combinedContent += `<jupyter_text>${text}`;
     }
-
   }
 
   return combinedContent;
 }
 
+
+async function getCellContentText(activeCell){
+  const result = await getChecked()
+
+  if(result == "openaiApiKey"){
+    return getCellContentTextRequiredForOpenAI(activeCell)
+  }else{
+    return getCellContentTextRequiredForBigCode(activeCell)
+  }
+}
+
+
+
 function extractTextFromCell(cell) {
   const codeMirrorLines = cell.querySelectorAll('.CodeMirror-code pre');
   const content = [];
 
   codeMirrorLines.forEach((line) => {
     content.push(line.textContent);
   });
+  const content_str = content.join('\n');
 
-  return content.join('\n');
+  return content_str;
 }
 
 
@@ -293,7 +335,9 @@ const startWaitingAnimation = (activeCell) => {
     const withAllCodeSpan = linesElement[lineIndex].querySelectorAll('span')
     withAllCodeSpan[withAllCodeSpan.length-1].appendChild(animationElement)
   }else{
-    currectLineSpanList[currectLineSpanList.length-1].insertAdjacentElement('afterend', animationElement);
+    // Insert new hint code in the last code span
+    const withAllCodeSpan = linesElement[lineIndex].childNodes
+    withAllCodeSpan[withAllCodeSpan.length - 1].insertAdjacentElement('afterend', animationElement);
   }
 
 
@@ -315,7 +359,7 @@ const startWaitingAnimation = (activeCell) => {
         break;
     }
 
-    animationElement.innerHTML = ' ' + animatedText + " time left: " + Math.floor(timeLeft-- / 3) + "s"
+    animationElement.innerText = ' ' + animatedText + " time left: " + Math.floor(timeLeft-- / 3) + "s"
 
     // request timeout
     if (timeLeft <= 0) {
@@ -373,15 +417,15 @@ if (document.querySelector('body.notebook_app')) {
       const activeCell = activeTextarea.parentElement.parentElement
 
       // Retrieve the content of the active cell 
-      const code = getCellContentText(activeCell);
+      const code = await getCellContentText(activeCell);
 
       if (!code) return;
 
       if (activeCell) {
         // Start Animation
         const [animationInterval, animationElement] = startWaitingAnimation(activeCell)
-
         isRequestInProgress = true
+
         const suggestion = await getCodeCompletion(code)
 
         if (suggestion) {

options.html

@@ -165,6 +165,8 @@
     <div class="input-key">
       <div>BigCode Url:</div>
       <div><input type="text" id="otherServiceUrl" /></div>
+      <div>Huggingface Api Key:</div>
+      <div><input type="text" id ="huggingfaceApiKey" /></div>
     </div>
 
     <div class="input-key input-hidden">

options.js

@@ -21,6 +21,7 @@ document.getElementById("save").addEventListener("click", () => {
   chrome.storage.sync.set({
     openaiApiKey: document.getElementById("apiKey").value,
     otherService: document.getElementById("otherServiceUrl").value,
+    huggingfaceApiKey: document.getElementById("huggingfaceApiKey").value,
     modelType: select.value,
     checked: getSelectedRadioValue()
   }, () => {
@@ -34,6 +35,12 @@ chrome.storage.sync.get("openaiApiKey", (data) => {
   }
 });
 
+chrome.storage.sync.get("huggingfaceApiKey", (data) => {
+  if (data.huggingfaceApiKey) {
+    document.getElementById("huggingfaceApiKey").value = data.huggingfaceApiKey;
+  }
+});
+
 chrome.storage.sync.get("otherService", (data) => {
   if (data.otherService) {
     document.getElementById("otherServiceUrl").value = data.otherService;