diff --git a/core/amber/src/main/resources/application.conf b/core/amber/src/main/resources/application.conf index 69e0c00188e..21f21a8af9a 100644 --- a/core/amber/src/main/resources/application.conf +++ b/core/amber/src/main/resources/application.conf @@ -91,3 +91,8 @@ region-plan-generator { enable-cost-based-region-plan-generator = false use-global-search = false } + +ai-assistant-server{ + assistant = "openai" + ai-service-key = "Put your OpenAI authentication key here" +} \ No newline at end of file diff --git a/core/amber/src/main/scala/edu/uci/ics/amber/engine/common/AmberConfig.scala b/core/amber/src/main/scala/edu/uci/ics/amber/engine/common/AmberConfig.scala index 6b1c04981ee..a5e3b8b4469 100644 --- a/core/amber/src/main/scala/edu/uci/ics/amber/engine/common/AmberConfig.scala +++ b/core/amber/src/main/scala/edu/uci/ics/amber/engine/common/AmberConfig.scala @@ -105,4 +105,6 @@ object AmberConfig { // JDBC configuration val jdbcConfig: Config = getConfSource.getConfig("jdbc") + // Python language server configuration + val aiAssistantConfig: Config = getConfSource.getConfig("ai-assistant-server") } diff --git a/core/amber/src/main/scala/edu/uci/ics/texera/web/TexeraWebApplication.scala b/core/amber/src/main/scala/edu/uci/ics/texera/web/TexeraWebApplication.scala index 99b8bfd19e7..e05e7e23fd3 100644 --- a/core/amber/src/main/scala/edu/uci/ics/texera/web/TexeraWebApplication.scala +++ b/core/amber/src/main/scala/edu/uci/ics/texera/web/TexeraWebApplication.scala @@ -258,6 +258,7 @@ class TexeraWebApplication environment.jersey.register(classOf[AdminExecutionResource]) environment.jersey.register(classOf[UserQuotaResource]) environment.jersey.register(classOf[UserDiscussionResource]) + environment.jersey.register(classOf[AiAssistantResource]) } /** diff --git a/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/AiAssistantManager.scala b/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/AiAssistantManager.scala new file mode 100644 index 00000000000..465adf46461 --- /dev/null +++ b/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/AiAssistantManager.scala @@ -0,0 +1,43 @@ +package edu.uci.ics.texera.web.resource.aiassistant +import edu.uci.ics.amber.engine.common.AmberConfig +import java.net.{HttpURLConnection, URL} +import java.util.logging.Logger + +object AiAssistantManager { + private val logger = Logger.getLogger(getClass.getName) + + private val aiAssistantConfig = AmberConfig.aiAssistantConfig + val assistantType: String = aiAssistantConfig.getString("assistant") + val accountKey: String = aiAssistantConfig.getString("ai-service-key") + val validAIAssistant: Boolean = assistantType match { + case "none" => + false + + case "openai" => + var isKeyValid: Boolean = false + var connection: HttpURLConnection = null + try { + val url = new URL("https://api.openai.com/v1/models") + connection = url.openConnection().asInstanceOf[HttpURLConnection] + connection.setRequestMethod("GET") + connection.setRequestProperty( + "Authorization", + s"Bearer ${accountKey.trim.replaceAll("^\"|\"$", "")}" + ) + val responseCode = connection.getResponseCode + isKeyValid = responseCode == 200 + } catch { + case e: Exception => + isKeyValid = false + logger.warning(s"Error validating OpenAI API key: ${e.getMessage}") + } finally { + if (connection != null) { + connection.disconnect() + } + } + isKeyValid + + case _ => + false + } +} diff --git a/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/AiAssistantResource.scala b/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/AiAssistantResource.scala new file mode 100644 index 00000000000..e42055bf034 --- /dev/null +++ b/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/AiAssistantResource.scala @@ -0,0 +1,132 @@ +package edu.uci.ics.texera.web.resource + +import edu.uci.ics.texera.web.auth.SessionUser +import edu.uci.ics.texera.web.resource.aiassistant.AiAssistantManager +import io.dropwizard.auth.Auth + +import javax.annotation.security.RolesAllowed +import javax.ws.rs._ + +import javax.ws.rs.core.Response +import java.util.Base64 +import scala.sys.process._ +import java.util.logging.Logger + +@Path("/aiassistant") +class AiAssistantResource { + private val logger = Logger.getLogger(classOf[AiAssistantResource].getName) + + final private lazy val isEnabled = AiAssistantManager.validAIAssistant + + @GET + @RolesAllowed(Array("REGULAR", "ADMIN")) + @Path("/isenabled") + def isAiAssistantEnabled: Boolean = isEnabled + + /** + * To get the type annotation suggestion from OpenAI + */ + @POST + @RolesAllowed(Array("REGULAR", "ADMIN")) + @Path("/getresult") + def getAiResponse(prompt: String, @Auth user: SessionUser): Response = { + val finalPrompt = prompt.replace("\\", "\\\\").replace("\"", "\\\"") + val requestBody = + s""" + |{ + | "model": "gpt-4", + | "messages": [{"role": "user", "content": "$finalPrompt"}], + | "max_tokens": 15 + |} + """.stripMargin + + try { + val url = new java.net.URL("https://api.openai.com/v1/chat/completions") + val connection = url.openConnection().asInstanceOf[java.net.HttpURLConnection] + connection.setRequestMethod("POST") + connection.setRequestProperty("Authorization", s"Bearer ${AiAssistantManager.accountKey}") + connection.setRequestProperty("Content-Type", "application/json") + connection.setDoOutput(true) + connection.getOutputStream.write(requestBody.getBytes("UTF-8")) + val responseCode = connection.getResponseCode + val responseStream = connection.getInputStream + val responseString = scala.io.Source.fromInputStream(responseStream).mkString + if (responseCode == 200) { + logger.info(s"Response from OpenAI API: $responseString") + } else { + logger.warning(s"Error response from OpenAI API: $responseString") + } + responseStream.close() + connection.disconnect() + Response.status(responseCode).entity(responseString).build() + } catch { + case e: Exception => + logger.warning(s"Exception occurred: ${e.getMessage}") + e.printStackTrace() + Response.status(Response.Status.INTERNAL_SERVER_ERROR).entity("Error occurred").build() + } + } + + @POST + @RolesAllowed(Array("REGULAR", "ADMIN")) + @Path("/getArgument") + def locateUnannotated(requestBody: AiAssistantResource.LocateUnannotatedRequest): Response = { + val selectedCode = requestBody.selectedCode + val startLine = requestBody.startLine + + val encodedCode = Base64.getEncoder.encodeToString(selectedCode.getBytes("UTF-8")) + try { + val pythonScriptPath = + "src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/python_abstract_syntax_tree.py" + val command = s"""python $pythonScriptPath "$encodedCode" $startLine""" + val result = command.!! + // Parse the string to the Json + val parsedJson = parseJson(result) + parsedJson match { + case Some(data: List[List[Any]]) => + val unannotatedArgs = data.map { + case List( + name: String, + startLine: Double, + startColumn: Double, + endLine: Double, + endColumn: Double + ) => + List(name, startLine.toInt, startColumn.toInt, endLine.toInt, endColumn.toInt) + } + logger.info(s"Unannotated arguments: $unannotatedArgs") + Response.ok(Map("result" -> unannotatedArgs)).build() + case _ => + Response.status(400).entity("Invalid JSON").build() + } + } catch { + case e: Exception => + Response.status(500).entity("Error with executing the python code").build() + } + } + + def parseJson(jsonString: String): Option[List[List[Any]]] = { + val cleanJson = jsonString.trim.drop(1).dropRight(1) + val rows = cleanJson.split("], \\[").toList + val parsedRows = rows.map { row => + val elements = row.replaceAll("[\\[\\]]", "").split(",").toList + if (elements.length == 5) { + val name = elements.head.trim.replace("\"", "") + val startLine = elements(1).trim.toDouble + val startColumn = elements(2).trim.toDouble + val endLine = elements(3).trim.toDouble + val endColumn = elements(4).trim.toDouble + List(name, startLine, startColumn, endLine, endColumn) + } else { + logger.warning("The Json format is wrong") + List.empty + } + } + val result = parsedRows.filter(_.nonEmpty) + Some(result) + } +} + +object AiAssistantResource { + case class LocateUnannotatedRequest(selectedCode: String, startLine: Int) +} diff --git a/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/python_abstract_syntax_tree.py b/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/python_abstract_syntax_tree.py new file mode 100644 index 00000000000..391f409e8bd --- /dev/null +++ b/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/python_abstract_syntax_tree.py @@ -0,0 +1,39 @@ +import ast +import json +import sys +import base64 + +class TypeAnnotationVisitor(ast.NodeVisitor): + def __init__(self, start_line_offset=0): + self.untyped_args = [] + # To calculate the correct start line + self.start_line_offset = start_line_offset + + def visit_FunctionDef(self, node): + for arg in node.args.args: + # Self is not an argument + if arg.arg == 'self': + continue + if not arg.annotation: + # +1 and -1 is to change ast line/col to monaco editor line/col for the range + start_line = arg.lineno + self.start_line_offset - 1 + start_col = arg.col_offset + 1 + end_line = start_line + end_col = start_col + len(arg.arg) + self.untyped_args.append([arg.arg, start_line, start_col, end_line, end_col]) + self.generic_visit(node) + +def find_untyped_variables(source_code, start_line): + tree = ast.parse(source_code) + visitor = TypeAnnotationVisitor(start_line_offset=start_line) + visitor.visit(tree) + return visitor.untyped_args + +if __name__ == "__main__": + # First argument is the code + encoded_code = sys.argv[1] + # Second argument is the start line + start_line = int(sys.argv[2]) + source_code = base64.b64decode(encoded_code).decode('utf-8') + untyped_variables = find_untyped_variables(source_code, start_line) + print(json.dumps(untyped_variables)) \ No newline at end of file diff --git a/core/gui/src/app/dashboard/service/user/ai-assistant/ai-assistant.service.ts b/core/gui/src/app/dashboard/service/user/ai-assistant/ai-assistant.service.ts new file mode 100644 index 00000000000..3a573eca507 --- /dev/null +++ b/core/gui/src/app/dashboard/service/user/ai-assistant/ai-assistant.service.ts @@ -0,0 +1,78 @@ +import { Injectable } from "@angular/core"; +import { firstValueFrom } from "rxjs"; +import { AppSettings } from "../../../../common/app-setting"; +import { HttpClient } from "@angular/common/http"; + +export const AI_ASSISTANT_API_BASE_URL = `${AppSettings.getApiEndpoint()}/aiassistant`; + +@Injectable({ + providedIn: "root", +}) +export class AiAssistantService { + constructor(private http: HttpClient) {} + + public checkAiAssistantEnabled(): Promise { + const apiUrl = `${AI_ASSISTANT_API_BASE_URL}/isenabled`; + return firstValueFrom(this.http.get(apiUrl)) + .then(response => (response !== undefined ? response : false)) + .catch(() => false); + } + + public getTypeAnnotations(code: string, lineNumber: number, allcode: string): Promise { + const prompt = ` + Your task is to analyze the given Python code and provide only the type annotation as stated in the instructions. + Instructions: + - The provided code will only be one of the 2 situations below: + - First situation: The input is not start with "def". If the provided code only contains variable, output the result in the format ":type". + - Second situation: The input is start with "def". If the provided code starts with "def" (a longer line than just a variable, indicative of a function or method), output the result in the format " -> type". + - The type should only be one word, such as "str", "int", etc. + + Examples: + - First situation: + - Provided code is "name", then the output may be : str + - Provided code is "age", then the output may be : int + - Provided code is "data", then the output may be : Tuple[int, str] + - Provided code is "new_user", then the output may be : User + - A special case: provided code is "self" and the context is something like "def __init__(self, username :str , age :int)", if the user requires the type annotation for the first parameter "self", then you should generate nothing. + - Second situation: (actual output depends on the complete code content) + - Provided code is "process_data(data: List[Tuple[int, str]], config: Dict[str, Union[int, str]])", then the output may be -> Optional[str] + - Provided code is "def add(a: int, b: int)", then the output may be -> int + + Counterexamples: + - Provided code is "def __init__(self, username: str, age: int)" and you generate the result: + The result is The provided code is "def __init__(self, username: str, age: int)", so it fits the second situation, which means the result should be in " -> type" format. However, the __init__ method in Python doesn't return anything or in other words, it implicitly returns None. Hence the correct type hint would be: -> None. + I don't want this result! The correct result you should generate is -> None for this counter case. + + Details: + - Provided code: ${code} + - Line number of the provided code in the complete code context: ${lineNumber} + - Complete code context: ${allcode} + + Important: (you must follow!!) + - For the first situation: you must return strictly according to the format ": type", without adding any extra characters. No need for an explanation, just the result : type is enough! + - For the second situation: you return strictly according to the format " -> type", without adding any extra characters. No need for an explanation, just the result -> type is enough! + `; + return firstValueFrom(this.http.post(`${AI_ASSISTANT_API_BASE_URL}/getresult`, { prompt })) + .then(response => { + console.log("Received response from backend:", response); + const result = response.choices[0].message.content.trim(); + return result; + }) + .catch(error => { + console.error("Request to backend failed:", error); + return ""; + }); + } + + public locateUnannotated(selectedCode: string, startLine: number) { + return firstValueFrom(this.http.post(`${AI_ASSISTANT_API_BASE_URL}/getArgument`, { selectedCode, startLine })) + .then(response => { + console.log("Received response from backend:", response); + return response.result; + }) + .catch(error => { + console.error("Request to backend failed:", error); + return []; + }); + } +} diff --git a/core/gui/src/app/workspace/component/code-editor-dialog/code-editor.component.html b/core/gui/src/app/workspace/component/code-editor-dialog/code-editor.component.html index e94d617541e..77560d23e41 100644 --- a/core/gui/src/app/workspace/component/code-editor-dialog/code-editor.component.html +++ b/core/gui/src/app/workspace/component/code-editor-dialog/code-editor.component.html @@ -29,3 +29,33 @@
+ +
+

Do you agree with the type annotation suggestion？

+ 
Adding annotation for code: {{ currentCode }}
+

Given suggestion: {{ currentSuggestion }}

+ + +
+ + diff --git a/core/gui/src/app/workspace/component/code-editor-dialog/code-editor.component.scss b/core/gui/src/app/workspace/component/code-editor-dialog/code-editor.component.scss index d4920dddeff..e62e77af5b7 100644 --- a/core/gui/src/app/workspace/component/code-editor-dialog/code-editor.component.scss +++ b/core/gui/src/app/workspace/component/code-editor-dialog/code-editor.component.scss @@ -49,3 +49,75 @@ left: -4px; top: -5px; } + +.annotation-suggestion { + position: absolute; + //bottom: 20px; + //right: 20px; + background: #222; + color: #fff; + padding: 20px; + border-radius: 8px; + box-shadow: 0px 0px 10px rgba(0, 0, 0, 0.5); + z-index: 1000; +} + +.annotation-suggestion button { + margin-right: 10px; +} + +.annotation-suggestion button.accept-button { + background-color: #28a745; + color: #000; + border: none; + padding: 10px 20px; + border-radius: 5px; + cursor: pointer; +} + +.annotation-suggestion button.accept-button:hover { + background-color: #218838; +} + +.annotation-suggestion button.decline-button { + background-color: #dc3545; + color: #000; + border: none; + padding: 10px 20px; + border-radius: 5px; + cursor: pointer; +} + +.annotation-suggestion button.decline-button:hover { + background-color: #c82333; +} + +.noAnnotationNeeded { + position: absolute; + background: #f5f5dc; + padding: 20px; + border-radius: 8px; + z-index: 1000; + top: 80px; + left: 80px; + max-width: 280px; + text-align: center; +} + +.noAnnotationNeeded.title { + margin-bottom: 15px; + font-size: 28px; +} + +.close-button-for-none { + background-color: #28a745; + color: #fff; + border: none; + padding: 6px 12px; + border-radius: 4px; + cursor: pointer; +} + +.close-button-for-none:hover { + background-color: #218838; +} diff --git a/core/gui/src/app/workspace/component/code-editor-dialog/code-editor.component.ts b/core/gui/src/app/workspace/component/code-editor-dialog/code-editor.component.ts index d0ed9343d0e..81d4866f30c 100644 --- a/core/gui/src/app/workspace/component/code-editor-dialog/code-editor.component.ts +++ b/core/gui/src/app/workspace/component/code-editor-dialog/code-editor.component.ts @@ -17,6 +17,7 @@ import { isUndefined } from "lodash"; import { CloseAction, ErrorAction } from "vscode-languageclient/lib/common/client.js"; import * as monaco from "monaco-editor/esm/vs/editor/editor.api.js"; import { FormControl } from "@angular/forms"; +import { AiAssistantService } from "../../../dashboard/service/user/ai-assistant/ai-assistant.service"; /** * CodeEditorComponent is the content of the dialogue invoked by CodeareaCustomTemplateComponent. @@ -46,6 +47,15 @@ export class CodeEditorComponent implements AfterViewInit, SafeStyle, OnDestroy public language: string = ""; public languageTitle: string = ""; + // Boolean to determine whether the suggestion UI should be shown + public showAnnotationSuggestion: boolean = false; + // The code selected by the user + public currentCode: string = ""; + // The result returned by the backend AI assistant + public currentSuggestion: string = ""; + // The range selected by the user + public currentRange: monaco.Range | null = null; + private generateLanguageTitle(language: string): string { return `${language.charAt(0).toUpperCase()}${language.slice(1)} UDF`; } @@ -63,7 +73,8 @@ export class CodeEditorComponent implements AfterViewInit, SafeStyle, OnDestroy private sanitizer: DomSanitizer, private workflowActionService: WorkflowActionService, private workflowVersionService: WorkflowVersionService, - public coeditorPresenceService: CoeditorPresenceService + public coeditorPresenceService: CoeditorPresenceService, + private aiAssistantService: AiAssistantService ) { const currentOperatorId = this.workflowActionService.getJointGraphWrapper().getCurrentHighlightedOperatorIDs()[0]; const operatorType = this.workflowActionService.getTexeraGraph().getOperator(currentOperatorId).operatorType; @@ -152,7 +163,7 @@ export class CodeEditorComponent implements AfterViewInit, SafeStyle, OnDestroy * Create a Monaco editor and connect it to MonacoBinding. * @private */ - private initMonaco() { + private async initMonaco() { const editor = monaco.editor.create(this.editorElement.nativeElement, { language: this.language, fontSize: 11, @@ -169,11 +180,307 @@ export class CodeEditorComponent implements AfterViewInit, SafeStyle, OnDestroy ); } this.editor = editor; + + // Check if the AI provider is "openai" + if (await this.aiAssistantService.checkAiAssistantEnabled()) { + // Add all needed modules for add type annotation + this.addAnnotationModule(editor); + + // "Add Type Annotation" Button + editor.addAction({ + id: "type-annotation-action", + label: "Add Type Annotation", + contextMenuGroupId: "1_modification", + contextMenuOrder: 1.0, + run: async ed => { + // User selected code(including range and content) + const selection = ed.getSelection(); + const model = ed.getModel(); + if (!model || !selection) { + return; + } + // All the code in Python UDF + const allcode = model.getValue(); + // Content of user selected code + const code = model.getValueInRange(selection); + // Start line of the selected code + const lineNumber = selection.startLineNumber; + await this.handleTypeAnnotation( + code, + selection, + ed as monaco.editor.IStandaloneCodeEditor, + lineNumber, + allcode + ); + }, + }); + + // "Add All Type Annotation" Button + editor.addAction({ + id: "all-type-annotation-action", + label: "Add All Type Annotations", + contextMenuGroupId: "1_modification", + contextMenuOrder: 1.1, + run: async ed => { + console.log("Add All Type Annotations action triggered"); + + const selection = ed.getSelection(); + const model = ed.getModel(); + if (!model || !selection) { + return; + } + + const selectedCode = model.getValueInRange(selection); + const allCode = model.getValue(); + // Locate the unannotated argument + const variablesWithoutAnnotations = await this.aiAssistantService.locateUnannotated( + selectedCode, + selection.startLineNumber + ); + + // If no argument need the type annotation, another UI will pop up + if (variablesWithoutAnnotations.length == 0) { + const popup = document.getElementById("noAnnotationNeeded"); + if (popup) { + popup.style.display = "block"; + const closeButton = document.getElementById("close-button-for-none"); + if (closeButton) { + closeButton.addEventListener("click", () => { + if (popup) { + popup.style.display = "none"; + } + }); + } + return; + } + } + + // Update range + let offset = 0; + let lastLine = null; + + for (let i = 0; i < variablesWithoutAnnotations.length; i++) { + const currVariable = variablesWithoutAnnotations[i]; + + // currVariable[0] is code + // currVariable[1] is start line + // currVariable[2] is start column + // currVariable[3] is end line + // currVariable[4] is end column + const variableCode = currVariable[0]; + const variableLineNumber = currVariable[1]; + + // If this variable is in the same line with the last variable, then the range(col) of this variable need to be updated + if (lastLine !== null && lastLine === variableLineNumber) { + offset += this.currentSuggestion.length; + } else { + // Reset to 0 if the variables are not in the same line + offset = 0; + } + const variableRange = new monaco.Range( + currVariable[1], + // Update the col after inserting last type annotation in the same line + currVariable[2] + offset, + currVariable[3], + currVariable[4] + offset + ); + + // Custom highlight for the current variable + const highlight = editor.createDecorationsCollection([ + { + range: variableRange, + options: { + hoverMessage: { value: "Argument without Annotation" }, + isWholeLine: false, + // Define in src/style.scss + className: "annotation-highlight", + }, + }, + ]); + + await this.handleTypeAnnotation( + variableCode, + variableRange, + ed as monaco.editor.IStandaloneCodeEditor, + variableLineNumber, + allCode + ); + // Clear the custom highlight + highlight.clear(); + // Update the lastLine + lastLine = variableLineNumber; + } + }, + }); + } + if (this.language == "python") { this.connectLanguageServer(); } } + private async handleTypeAnnotation( + code: string, + range: monaco.Range, + editor: monaco.editor.IStandaloneCodeEditor, + lineNumber: number, + allcode: string + ): Promise { + return new Promise(resolve => { + this.aiAssistantService.getTypeAnnotations(code, lineNumber, allcode).then(typeAnnotations => { + console.log("The result from OpenAI is", typeAnnotations); + + let acceptButton: HTMLButtonElement | null = null; + let declineButton: HTMLButtonElement | null = null; + + this.currentCode = code; + this.currentSuggestion = typeAnnotations; + this.currentRange = range; + this.showAnnotationSuggestion = true; + + // Let the suggestion pop up next to the selected code + setTimeout(() => { + const position = editor.getScrolledVisiblePosition(range.getStartPosition()); + const popupElement = document.querySelector(".annotation-suggestion") as HTMLElement; + + if (popupElement && position) { + popupElement.style.top = `${position.top + 100}px`; + popupElement.style.left = `${position.left + 100}px`; + } + + // Make sure the user click the button + const cleanup = () => { + console.log("Cleaning up and resolving..."); + if (acceptButton) acceptButton.removeEventListener("click", acceptListener); + if (declineButton) declineButton.removeEventListener("click", declineListener); + this.showAnnotationSuggestion = false; + resolve(); + console.log("Resolved!"); + }; + + const acceptListener = () => { + this.acceptCurrentAnnotation(); + cleanup(); + }; + + const declineListener = () => { + cleanup(); + }; + acceptButton = document.querySelector(".accept-button") as HTMLButtonElement; + declineButton = document.querySelector(".decline-button") as HTMLButtonElement; + + if (acceptButton && declineButton) { + console.log("Buttons found, adding event listeners"); + //clean the old one for the "add all type annotation" + acceptButton.removeEventListener("click", acceptListener); + declineButton.removeEventListener("click", declineListener); + + acceptButton.addEventListener("click", acceptListener, { once: true }); + declineButton.addEventListener("click", declineListener, { once: true }); + } else { + console.error("Buttons not found!"); + } + }, 0); + }); + }); + } + + // Called when the user clicks the "accept" button + public acceptCurrentAnnotation(): void { + // Avoid accidental calls + if (!this.showAnnotationSuggestion || !this.currentRange || !this.currentSuggestion) { + return; + } + + if (this.currentRange && this.currentSuggestion) { + const selection = new monaco.Selection( + this.currentRange.startLineNumber, + this.currentRange.startColumn, + this.currentRange.endLineNumber, + this.currentRange.endColumn + ); + this.insertTypeAnnotations(this.editor, selection, this.currentSuggestion); + } + // close the UI after adding the annotation + this.showAnnotationSuggestion = false; + } + + // Called when the user clicks the "decline" button + public rejectCurrentAnnotation(): void { + // Do nothing except for closing the UI + this.showAnnotationSuggestion = false; + } + + // Add the type annotation into monaco editor + private insertTypeAnnotations( + editor: monaco.editor.IStandaloneCodeEditor, + selection: monaco.Selection, + annotations: string + ) { + const endLineNumber = selection.endLineNumber; + const endColumn = selection.endColumn; + const range = new monaco.Range( + // Insert the content to the end of the selected code + endLineNumber, + endColumn, + endLineNumber, + endColumn + ); + const text = `${annotations}`; + const op = { + range: range, + text: text, + forceMoveMarkers: true, + }; + editor.executeEdits("add annotation", [op]); + } + + // Add all necessary modules for type annotation at the first line of the Python UDF + private addAnnotationModule(editor: monaco.editor.IStandaloneCodeEditor) { + const model = editor.getModel(); + if (!model) { + return; + } + const allCode = model.getValue(); + const typingImports = [ + "Any", + "Awaitable", + "Callable", + "Coroutine", + "Dict", + "FrozenSet", + "Generator", + "Generic", + "Iterable", + "Iterator", + "List", + "Mapping", + "Optional", + "Sequence", + "Set", + "Tuple", + "Type", + "TypeVar", + "Union", + "Deque", + "NamedTuple", + "TypedDict", + "Protocol", + "Literal", + "NewType", + "NoReturn", + ]; + const importStatement = `from typing import (\n ${typingImports.join(",\n ")}\n)`; + if (!allCode.includes(importStatement)) { + const importOp = { + // Add the module at the first line + range: new monaco.Range(1, 1, 1, 1), + text: `${importStatement}\n\n`, + }; + editor.executeEdits("add module", [importOp]); + } + } + private connectLanguageServer() { if (this.languageServerSocket === undefined) { this.languageServerSocket = new WebSocket(getWebsocketUrl("/python-language-server", "3000")); @@ -197,6 +504,16 @@ export class CodeEditorComponent implements AfterViewInit, SafeStyle, OnDestroy reader.onClose(() => languageClient.stop()); } }; + // Make sure that the Pyright will be reconnected if the user refresh the website + this.languageServerSocket.onclose = () => { + console.log("WebSocket connection will be reconnect"); + setTimeout(() => { + this.connectLanguageServer(); + }, 3000); + }; + this.languageServerSocket.onerror = error => { + console.error("WebSocket error:", error); + }; } } diff --git a/core/gui/src/styles.scss b/core/gui/src/styles.scss index d413e250da0..6c2ac14306b 100644 --- a/core/gui/src/styles.scss +++ b/core/gui/src/styles.scss @@ -72,3 +72,7 @@ hr { .ant-tabs-tabpane { padding-right: 24px; } + +.annotation-highlight { + background-color: #6a5acd; +}