From 4743a649cd1f22599cb2e20b3282ef2a69bf2050 Mon Sep 17 00:00:00 2001
From: "Minchong(Brian) Wu" <2638932112@qq.com>
Date: Sat, 21 Sep 2024 02:48:45 -0400
Subject: [PATCH 01/12] initial commit

---
 .../aiassistant/AiAssistantResource.scala     |  42 +++++++
 .../python_abstract_syntax_tree.py            |  39 +++++++
 .../code-editor.component.ts                  | 104 +++++++++++++++++-
 .../ai-assistant/ai-assistant.service.ts      |  52 +++++++++
 core/gui/src/styles.scss                      |   4 +
 5 files changed, 239 insertions(+), 2 deletions(-)
 create mode 100644 core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/python_abstract_syntax_tree.py

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
index f784165dc65..6cbdd63878f 100644
--- 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
@@ -9,8 +9,17 @@ import javax.ws.rs.Consumes
 import javax.ws.rs.core.MediaType
 import play.api.libs.json.Json
 import kong.unirest.Unirest
+import java.util.Base64
+import scala.sys.process._
+import play.api.libs.json._
+import scala.util.{Try, Success, Failure}
 
 case class AIAssistantRequest(code: String, lineNumber: Int, allcode: String)
+case class LocateUnannotatedRequest(selectedCode: String, startLine: Int)
+case class UnannotatedArgument(name: String, startLine: Int, startColumn: Int, endLine: Int, endColumn: Int)
+object UnannotatedArgument {
+  implicit val format: Format[UnannotatedArgument] = Json.format[UnannotatedArgument]
+}
 
 @Path("/aiassistant")
 class AIAssistantResource {
@@ -86,4 +95,37 @@ class AIAssistantResource {
        |- 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!
      """.stripMargin
   }
+
+  @POST
+  @RolesAllowed(Array("REGULAR", "ADMIN"))
+  @Path("/getArgument")
+  @Consumes(Array(MediaType.APPLICATION_JSON))
+  def locateUnannotated(request: LocateUnannotatedRequest, @Auth user: SessionUser): Response = {
+    val encodedCode = Base64.getEncoder.encodeToString(request.selectedCode.getBytes("UTF-8"))
+    val pythonScriptPath = "src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/python_abstract_syntax_tree.py"
+
+    try {
+      val command = s"""python $pythonScriptPath "$encodedCode" ${request.startLine}"""
+      val result = command.!!
+      val parsedResult = parseJsonResult(result)
+      Response.ok(Json.obj("result" -> Json.toJson(parsedResult))).build()
+    } catch {
+      case e: Exception =>
+        Response.status(500).entity("Error executing the Python code").build()
+    }
+  }
+
+  private def parseJsonResult(jsonString: String): List[UnannotatedArgument] = {
+    val jsonValue = Json.parse(jsonString)
+    jsonValue.as[List[List[JsValue]]].map {
+      case List(name: JsString, startLine: JsNumber, startColumn: JsNumber, endLine: JsNumber, endColumn: JsNumber) =>
+        UnannotatedArgument(
+          name.value,
+          startLine.value.toInt,
+          startColumn.value.toInt,
+          endLine.value.toInt,
+          endColumn.value.toInt
+        )
+    }
+  }
 }
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/workspace/component/code-editor-dialog/code-editor.component.ts b/core/gui/src/app/workspace/component/code-editor-dialog/code-editor.component.ts
index c772ad7786c..679afc37c18 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
@@ -4,7 +4,7 @@ import { WorkflowActionService } from "../../service/workflow-graph/model/workfl
 import { WorkflowVersionService } from "../../../dashboard/service/user/workflow-version/workflow-version.service";
 import { YText } from "yjs/dist/src/types/YText";
 import { MonacoBinding } from "y-monaco";
-import { Subject } from "rxjs";
+import { Subject, take } from "rxjs";
 import { takeUntil } from "rxjs/operators";
 import { MonacoLanguageClient } from "monaco-languageclient";
 import { toSocket, WebSocketMessageReader, WebSocketMessageWriter } from "vscode-ws-jsonrpc";
@@ -59,6 +59,9 @@ export class CodeEditorComponent implements AfterViewInit, SafeStyle, OnDestroy
   public currentRange: monaco.Range | undefined;
   public suggestionTop: number = 0;
   public suggestionLeft: number = 0;
+  // For "Add All Type Annotation" to show the UI individually
+  private userResponseSubject?: Subject<void>;
+  private isMultipleVariables: boolean = false;
 
   private generateLanguageTitle(language: string): string {
     return `${language.charAt(0).toUpperCase()}${language.slice(1)} UDF`;
@@ -221,7 +224,94 @@ export class CodeEditorComponent implements AfterViewInit, SafeStyle, OnDestroy
               },
             });
           }
-        },
+
+          // "Add All Type Annotation" Button
+          editor.addAction({
+            id: "all-type-annotation-action",
+            label: "Add All Type Annotations",
+            contextMenuGroupId: "1_modification",
+            contextMenuOrder: 1.1,
+            run: ed => {
+              const selection = ed.getSelection();
+              const model = ed.getModel();
+              if (!model || !selection) {
+                return;
+              }
+
+              const selectedCode = model.getValueInRange(selection);
+              const allCode = model.getValue();
+
+              this.aiAssistantService.locateUnannotated(selectedCode, selection.startLineNumber)
+                .pipe(takeUntil(this.workflowVersionStreamSubject))
+                .subscribe(variablesWithoutAnnotations => {
+                  // If no unannotated variable, then do nothing.
+                  if (variablesWithoutAnnotations.length == 0) {
+                    return;
+                  }
+
+                  let offset = 0;
+                  let lastLine: number | undefined;
+
+                  this.isMultipleVariables = true;
+                  this.userResponseSubject = new Subject<void>();
+
+                  const processNextVariable = (index: number) => {
+                    if (index >= variablesWithoutAnnotations.length) {
+                      this.isMultipleVariables = false;
+                      this.userResponseSubject = undefined;
+                      return;
+                    }
+
+                    const currVariable = variablesWithoutAnnotations[index];
+
+                    const variableCode = currVariable.name;
+                    const variableLineNumber = currVariable.startLine;
+
+                    // Update range
+                    if (lastLine !== undefined && lastLine === variableLineNumber) {
+                      offset += this.currentSuggestion.length;
+                    } else {
+                      offset = 0;
+                    }
+
+                    const variableRange = new monaco.Range(
+                      currVariable.startLine,
+                      currVariable.startColumn + offset,
+                      currVariable.endLine,
+                      currVariable.endColumn + offset
+                    );
+
+                    const highlight = editor.createDecorationsCollection([{
+                      range: variableRange,
+                      options: {
+                        hoverMessage: { value: "Argument without Annotation" },
+                        isWholeLine: false,
+                        className: "annotation-highlight",
+                      },
+                    }]);
+
+                    this.handleTypeAnnotation(variableCode, variableRange, ed as monaco.editor.IStandaloneCodeEditor, variableLineNumber, allCode);
+
+                    lastLine = variableLineNumber;
+
+                    // Make sure the currVariable will not go to the next one until the user click the accept/decline button
+                    if (this.userResponseSubject !== undefined) {
+                      const userResponseSubject = this.userResponseSubject;
+                      // Only take one response (accept/decline)
+                      const subscription = userResponseSubject.pipe(take(1))
+                        .pipe(takeUntil(this.workflowVersionStreamSubject))
+                        .subscribe(() => {
+                          highlight.clear();
+                          subscription.unsubscribe();
+                          processNextVariable(index + 1);
+                        });
+                    }
+                  };
+                  processNextVariable(0);
+                });
+            },
+          });
+        }
       });
     if (this.language == "python") {
       this.connectLanguageServer();
@@ -285,6 +375,11 @@ export class CodeEditorComponent implements AfterViewInit, SafeStyle, OnDestroy
         this.currentRange.endColumn
       );
       this.insertTypeAnnotations(this.editor, selection, this.currentSuggestion);
+
+      // Only for "Add All Type Annotation"
+      if (this.isMultipleVariables && this.userResponseSubject) {
+        this.userResponseSubject.next();
+      }
     }
     // close the UI after adding the annotation
     this.showAnnotationSuggestion = false;
@@ -296,6 +391,11 @@ export class CodeEditorComponent implements AfterViewInit, SafeStyle, OnDestroy
     this.showAnnotationSuggestion = false;
     this.currentCode = "";
     this.currentSuggestion = "";
+
+    // Only for "Add All Type Annotation"
+    if (this.isMultipleVariables && this.userResponseSubject) {
+      this.userResponseSubject.next();
+    }
   }
 
   // Add the type annotation into monaco editor
diff --git a/core/gui/src/app/workspace/service/ai-assistant/ai-assistant.service.ts b/core/gui/src/app/workspace/service/ai-assistant/ai-assistant.service.ts
index 0463b6af309..58ee8205794 100644
--- a/core/gui/src/app/workspace/service/ai-assistant/ai-assistant.service.ts
+++ b/core/gui/src/app/workspace/service/ai-assistant/ai-assistant.service.ts
@@ -13,6 +13,32 @@ export type TypeAnnotationResponse = {
   }>;
 };
 
+export interface UnannotatedArgument {
+  name: string;
+  startLine: number;
+  startColumn: number;
+  endLine: number;
+  endColumn: number;
+}
+
+interface UnannotatedArgumentItem {
+  underlying: {
+    name: { value: string };
+    startLine: { value: number };
+    startColumn: { value: number };
+    endLine: { value: number };
+    endColumn: { value: number };
+  };
+}
+
+interface UnannotatedArgumentResponse {
+  underlying: {
+    result: {
+      value: UnannotatedArgumentItem[];
+    };
+  };
+}
+
 // Define AI model type
 export const AI_ASSISTANT_API_BASE_URL = `${AppSettings.getApiEndpoint()}/aiassistant`;
 export const AI_MODEL = {
@@ -66,4 +92,30 @@ export class AIAssistantService {
     const requestBody = { code, lineNumber, allcode };
     return this.http.post<TypeAnnotationResponse>(`${AI_ASSISTANT_API_BASE_URL}/annotationresult`, requestBody, {});
   }
+
+  public locateUnannotated(selectedCode: string, startLine: number): Observable<UnannotatedArgument[]> {
+    const requestBody = { selectedCode, startLine };
+
+    return this.http.post<UnannotatedArgumentResponse>(`${AI_ASSISTANT_API_BASE_URL}/getArgument`, requestBody)
+      .pipe(
+        map(response => {
+          if (response) {
+            return response.underlying.result.value.map((item: UnannotatedArgumentItem): UnannotatedArgument => ({
+              name: item.underlying.name.value,
+              startLine: item.underlying.startLine.value,
+              startColumn: item.underlying.startColumn.value,
+              endLine: item.underlying.endLine.value,
+              endColumn: item.underlying.endColumn.value
+            }));
+          } else {
+            console.error("Unexpected response format:", response);
+            return [];
+          }
+        }),
+        catchError((error: unknown) => {
+          console.error("Request to backend failed:", error);
+          return of([]);
+        })
+      );
+  }
 }
diff --git a/core/gui/src/styles.scss b/core/gui/src/styles.scss
index 8f1e3f9e7eb..a8109d9439f 100644
--- a/core/gui/src/styles.scss
+++ b/core/gui/src/styles.scss
@@ -76,3 +76,7 @@ hr {
 body {
   overflow: hidden;
 }
+
+.annotation-highlight {
+  background-color: #6a5acd;
+}

From b31731cbc550a383f9f2574be092104002052e48 Mon Sep 17 00:00:00 2001
From: "Minchong(Brian) Wu" <2638932112@qq.com>
Date: Sat, 21 Sep 2024 02:53:37 -0400
Subject: [PATCH 02/12] fix fmt

---
 .../aiassistant/AiAssistantResource.scala     | 20 ++++++++---
 .../code-editor.component.ts                  | 34 +++++++++++-------
 .../ai-assistant/ai-assistant.service.ts      | 35 ++++++++++---------
 3 files changed, 56 insertions(+), 33 deletions(-)

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
index 6cbdd63878f..21287cc6d68 100644
--- 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
@@ -12,11 +12,16 @@ import kong.unirest.Unirest
 import java.util.Base64
 import scala.sys.process._
 import play.api.libs.json._
-import scala.util.{Try, Success, Failure}
 
 case class AIAssistantRequest(code: String, lineNumber: Int, allcode: String)
 case class LocateUnannotatedRequest(selectedCode: String, startLine: Int)
-case class UnannotatedArgument(name: String, startLine: Int, startColumn: Int, endLine: Int, endColumn: Int)
+case class UnannotatedArgument(
+    name: String,
+    startLine: Int,
+    startColumn: Int,
+    endLine: Int,
+    endColumn: Int
+)
 object UnannotatedArgument {
   implicit val format: Format[UnannotatedArgument] = Json.format[UnannotatedArgument]
 }
@@ -102,7 +107,8 @@ class AIAssistantResource {
   @Consumes(Array(MediaType.APPLICATION_JSON))
   def locateUnannotated(request: LocateUnannotatedRequest, @Auth user: SessionUser): Response = {
     val encodedCode = Base64.getEncoder.encodeToString(request.selectedCode.getBytes("UTF-8"))
-    val pythonScriptPath = "src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/python_abstract_syntax_tree.py"
+    val pythonScriptPath =
+      "src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/python_abstract_syntax_tree.py"
 
     try {
       val command = s"""python $pythonScriptPath "$encodedCode" ${request.startLine}"""
@@ -118,7 +124,13 @@ class AIAssistantResource {
   private def parseJsonResult(jsonString: String): List[UnannotatedArgument] = {
     val jsonValue = Json.parse(jsonString)
     jsonValue.as[List[List[JsValue]]].map {
-      case List(name: JsString, startLine: JsNumber, startColumn: JsNumber, endLine: JsNumber, endColumn: JsNumber) =>
+      case List(
+            name: JsString,
+            startLine: JsNumber,
+            startColumn: JsNumber,
+            endLine: JsNumber,
+            endColumn: JsNumber
+          ) =>
         UnannotatedArgument(
           name.value,
           startLine.value.toInt,
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 679afc37c18..b8193bb1844 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
@@ -241,7 +241,8 @@ export class CodeEditorComponent implements AfterViewInit, SafeStyle, OnDestroy
               const selectedCode = model.getValueInRange(selection);
               const allCode = model.getValue();
 
-              this.aiAssistantService.locateUnannotated(selectedCode, selection.startLineNumber)
+              this.aiAssistantService
+                .locateUnannotated(selectedCode, selection.startLineNumber)
                 .pipe(takeUntil(this.workflowVersionStreamSubject))
                 .subscribe(variablesWithoutAnnotations => {
                   // If no unannotated variable, then do nothing.
@@ -281,16 +282,24 @@ export class CodeEditorComponent implements AfterViewInit, SafeStyle, OnDestroy
                       currVariable.endColumn + offset
                     );
 
-                    const highlight = editor.createDecorationsCollection([{
-                      range: variableRange,
-                      options: {
-                        hoverMessage: { value: "Argument without Annotation" },
-                        isWholeLine: false,
-                        className: "annotation-highlight",
+                    const highlight = editor.createDecorationsCollection([
+                      {
+                        range: variableRange,
+                        options: {
+                          hoverMessage: { value: "Argument without Annotation" },
+                          isWholeLine: false,
+                          className: "annotation-highlight",
+                        },
                       },
-                    }]);
-
-                    this.handleTypeAnnotation(variableCode, variableRange, ed as monaco.editor.IStandaloneCodeEditor, variableLineNumber, allCode);
+                    ]);
+
+                    this.handleTypeAnnotation(
+                      variableCode,
+                      variableRange,
+                      ed as monaco.editor.IStandaloneCodeEditor,
+                      variableLineNumber,
+                      allCode
+                    );
 
                     lastLine = variableLineNumber;
 
@@ -298,7 +307,8 @@ export class CodeEditorComponent implements AfterViewInit, SafeStyle, OnDestroy
                     if (this.userResponseSubject !== undefined) {
                       const userResponseSubject = this.userResponseSubject;
                       // Only take one response (accept/decline)
-                      const subscription = userResponseSubject.pipe(take(1))
+                      const subscription = userResponseSubject
+                        .pipe(take(1))
                         .pipe(takeUntil(this.workflowVersionStreamSubject))
                         .subscribe(() => {
                           highlight.clear();
@@ -311,7 +321,7 @@ export class CodeEditorComponent implements AfterViewInit, SafeStyle, OnDestroy
                 });
             },
           });
-        }
+        },
       });
     if (this.language == "python") {
       this.connectLanguageServer();
diff --git a/core/gui/src/app/workspace/service/ai-assistant/ai-assistant.service.ts b/core/gui/src/app/workspace/service/ai-assistant/ai-assistant.service.ts
index 58ee8205794..7eea438fd9c 100644
--- a/core/gui/src/app/workspace/service/ai-assistant/ai-assistant.service.ts
+++ b/core/gui/src/app/workspace/service/ai-assistant/ai-assistant.service.ts
@@ -96,26 +96,27 @@ export class AIAssistantService {
   public locateUnannotated(selectedCode: string, startLine: number): Observable<UnannotatedArgument[]> {
     const requestBody = { selectedCode, startLine };
 
-    return this.http.post<UnannotatedArgumentResponse>(`${AI_ASSISTANT_API_BASE_URL}/getArgument`, requestBody)
-      .pipe(
-        map(response => {
-          if (response) {
-            return response.underlying.result.value.map((item: UnannotatedArgumentItem): UnannotatedArgument => ({
+    return this.http.post<UnannotatedArgumentResponse>(`${AI_ASSISTANT_API_BASE_URL}/getArgument`, requestBody).pipe(
+      map(response => {
+        if (response) {
+          return response.underlying.result.value.map(
+            (item: UnannotatedArgumentItem): UnannotatedArgument => ({
               name: item.underlying.name.value,
               startLine: item.underlying.startLine.value,
               startColumn: item.underlying.startColumn.value,
               endLine: item.underlying.endLine.value,
-              endColumn: item.underlying.endColumn.value
-            }));
-          } else {
-            console.error("Unexpected response format:", response);
-            return [];
-          }
-        }),
-        catchError((error: unknown) => {
-          console.error("Request to backend failed:", error);
-          return of([]);
-        })
-      );
+              endColumn: item.underlying.endColumn.value,
+            })
+          );
+        } else {
+          console.error("Unexpected response format:", response);
+          return [];
+        }
+      }),
+      catchError((error: unknown) => {
+        console.error("Request to backend failed:", error);
+        return of([]);
+      })
+    );
   }
 }

From 95d09b492101691f6c8d41133a5dde810a7b0eca Mon Sep 17 00:00:00 2001
From: IamtherealBrian <2638932112@qq.com>
Date: Mon, 23 Sep 2024 23:44:48 -0400
Subject: [PATCH 03/12] fix part comments

---
 .../aiassistant/AiAssistantResource.scala     |  7 +++--
 ...tax_tree.py => type_annotation_visitor.py} |  0
 .../ai-assistant/ai-assistant.service.ts      | 31 ++++++++++---------
 3 files changed, 20 insertions(+), 18 deletions(-)
 rename core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/{python_abstract_syntax_tree.py => type_annotation_visitor.py} (100%)

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
index 21287cc6d68..3d17b59abd1 100644
--- 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
@@ -12,6 +12,7 @@ import kong.unirest.Unirest
 import java.util.Base64
 import scala.sys.process._
 import play.api.libs.json._
+import java.nio.file.Paths
 
 case class AIAssistantRequest(code: String, lineNumber: Int, allcode: String)
 case class LocateUnannotatedRequest(selectedCode: String, startLine: Int)
@@ -103,12 +104,12 @@ class AIAssistantResource {
 
   @POST
   @RolesAllowed(Array("REGULAR", "ADMIN"))
-  @Path("/getArgument")
+  @Path("/annotate-argument")
   @Consumes(Array(MediaType.APPLICATION_JSON))
   def locateUnannotated(request: LocateUnannotatedRequest, @Auth user: SessionUser): Response = {
     val encodedCode = Base64.getEncoder.encodeToString(request.selectedCode.getBytes("UTF-8"))
     val pythonScriptPath =
-      "src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/python_abstract_syntax_tree.py"
+      Paths.get("src", "main", "scala", "edu", "uci", "ics", "texera", "web", "resource", "aiassistant", "type_annotation_visitor.py").toString
 
     try {
       val command = s"""python $pythonScriptPath "$encodedCode" ${request.startLine}"""
@@ -117,7 +118,7 @@ class AIAssistantResource {
       Response.ok(Json.obj("result" -> Json.toJson(parsedResult))).build()
     } catch {
       case e: Exception =>
-        Response.status(500).entity("Error executing the Python code").build()
+        Response.status(500).entity(s"Error executing the Python code: ${e.getMessage}").build()
     }
   }
 
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/type_annotation_visitor.py
similarity index 100%
rename from core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/python_abstract_syntax_tree.py
rename to core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/type_annotation_visitor.py
diff --git a/core/gui/src/app/workspace/service/ai-assistant/ai-assistant.service.ts b/core/gui/src/app/workspace/service/ai-assistant/ai-assistant.service.ts
index 7eea438fd9c..313fcd60f3c 100644
--- a/core/gui/src/app/workspace/service/ai-assistant/ai-assistant.service.ts
+++ b/core/gui/src/app/workspace/service/ai-assistant/ai-assistant.service.ts
@@ -14,31 +14,32 @@ export type TypeAnnotationResponse = {
 };
 
 export interface UnannotatedArgument {
-  name: string;
-  startLine: number;
-  startColumn: number;
-  endLine: number;
-  endColumn: number;
+  readonly name: string;
+  readonly startLine: number;
+  readonly startColumn: number;
+  readonly endLine: number;
+  readonly endColumn: number;
 }
 
 interface UnannotatedArgumentItem {
-  underlying: {
-    name: { value: string };
-    startLine: { value: number };
-    startColumn: { value: number };
-    endLine: { value: number };
-    endColumn: { value: number };
+  readonly underlying: {
+    readonly name: { readonly value: string };
+    readonly startLine: { readonly value: number };
+    readonly startColumn: { readonly value: number };
+    readonly endLine: { readonly value: number };
+    readonly endColumn: { readonly value: number };
   };
 }
 
 interface UnannotatedArgumentResponse {
-  underlying: {
-    result: {
-      value: UnannotatedArgumentItem[];
+  readonly underlying: {
+    readonly result: {
+      readonly value: ReadonlyArray<UnannotatedArgumentItem>;
     };
   };
 }
 
+
 // Define AI model type
 export const AI_ASSISTANT_API_BASE_URL = `${AppSettings.getApiEndpoint()}/aiassistant`;
 export const AI_MODEL = {
@@ -96,7 +97,7 @@ export class AIAssistantService {
   public locateUnannotated(selectedCode: string, startLine: number): Observable<UnannotatedArgument[]> {
     const requestBody = { selectedCode, startLine };
 
-    return this.http.post<UnannotatedArgumentResponse>(`${AI_ASSISTANT_API_BASE_URL}/getArgument`, requestBody).pipe(
+    return this.http.post<UnannotatedArgumentResponse>(`${AI_ASSISTANT_API_BASE_URL}/annotate-argument`, requestBody).pipe(
       map(response => {
         if (response) {
           return response.underlying.result.value.map(

From 24a74cab4bdc640da00a24583cd57ae658a49a11 Mon Sep 17 00:00:00 2001
From: IamtherealBrian <2638932112@qq.com>
Date: Tue, 24 Sep 2024 01:41:11 -0400
Subject: [PATCH 04/12] fix part comment

---
 .../aiassistant/AiAssistantResource.scala     | 32 +++++++------------
 .../code-editor.component.ts                  |  9 ++++--
 2 files changed, 17 insertions(+), 24 deletions(-)

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
index 3d17b59abd1..613a57729b7 100644
--- 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
@@ -13,6 +13,8 @@ import java.util.Base64
 import scala.sys.process._
 import play.api.libs.json._
 import java.nio.file.Paths
+import com.fasterxml.jackson.databind.ObjectMapper
+import com.fasterxml.jackson.module.scala.DefaultScalaModule
 
 case class AIAssistantRequest(code: String, lineNumber: Int, allcode: String)
 case class LocateUnannotatedRequest(selectedCode: String, startLine: Int)
@@ -29,6 +31,8 @@ object UnannotatedArgument {
 
 @Path("/aiassistant")
 class AIAssistantResource {
+  val objectMapper = new ObjectMapper()
+  objectMapper.registerModule(DefaultScalaModule)
   final private lazy val isEnabled = AiAssistantManager.validAIAssistant
   @GET
   @RolesAllowed(Array("REGULAR", "ADMIN"))
@@ -114,31 +118,17 @@ class AIAssistantResource {
     try {
       val command = s"""python $pythonScriptPath "$encodedCode" ${request.startLine}"""
       val result = command.!!
-      val parsedResult = parseJsonResult(result)
+    val parsedResult = objectMapper.readValue(result, classOf[List[List[Any]]]).map {
+    case List(name: String, startLine: Int, startColumn: Int, endLine: Int, endColumn: Int) =>
+      UnannotatedArgument(name, startLine, startColumn, endLine, endColumn)
+    case _ =>
+      throw new RuntimeException("Unexpected format in Python script result")
+    }
       Response.ok(Json.obj("result" -> Json.toJson(parsedResult))).build()
     } catch {
       case e: Exception =>
+        e.printStackTrace()
         Response.status(500).entity(s"Error executing the Python code: ${e.getMessage}").build()
     }
   }
-
-  private def parseJsonResult(jsonString: String): List[UnannotatedArgument] = {
-    val jsonValue = Json.parse(jsonString)
-    jsonValue.as[List[List[JsValue]]].map {
-      case List(
-            name: JsString,
-            startLine: JsNumber,
-            startColumn: JsNumber,
-            endLine: JsNumber,
-            endColumn: JsNumber
-          ) =>
-        UnannotatedArgument(
-          name.value,
-          startLine.value.toInt,
-          startColumn.value.toInt,
-          endLine.value.toInt,
-          endColumn.value.toInt
-        )
-    }
-  }
 }
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 b8193bb1844..8cbaf356c4d 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
@@ -62,6 +62,8 @@ export class CodeEditorComponent implements AfterViewInit, SafeStyle, OnDestroy
   // For "Add All Type Annotation" to show the UI individually
   private userResponseSubject?: Subject<void>;
   private isMultipleVariables: boolean = false;
+  private componentDestroy = new Subject<void>;
+
 
   private generateLanguageTitle(language: string): string {
     return `${language.charAt(0).toUpperCase()}${language.slice(1)} UDF`;
@@ -243,7 +245,7 @@ export class CodeEditorComponent implements AfterViewInit, SafeStyle, OnDestroy
 
               this.aiAssistantService
                 .locateUnannotated(selectedCode, selection.startLineNumber)
-                .pipe(takeUntil(this.workflowVersionStreamSubject))
+                .pipe(takeUntil(this.componentDestroy))
                 .subscribe(variablesWithoutAnnotations => {
                   // If no unannotated variable, then do nothing.
                   if (variablesWithoutAnnotations.length == 0) {
@@ -309,7 +311,7 @@ export class CodeEditorComponent implements AfterViewInit, SafeStyle, OnDestroy
                       // Only take one response (accept/decline)
                       const subscription = userResponseSubject
                         .pipe(take(1))
-                        .pipe(takeUntil(this.workflowVersionStreamSubject))
+                        .pipe(takeUntil(this.componentDestroy))
                         .subscribe(() => {
                           highlight.clear();
                           subscription.unsubscribe();
@@ -328,6 +330,7 @@ export class CodeEditorComponent implements AfterViewInit, SafeStyle, OnDestroy
     }
   }
 
+
   private handleTypeAnnotation(
     code: string,
     range: monaco.Range,
@@ -337,7 +340,7 @@ export class CodeEditorComponent implements AfterViewInit, SafeStyle, OnDestroy
   ): void {
     this.aiAssistantService
       .getTypeAnnotations(code, lineNumber, allcode)
-      .pipe(takeUntil(this.workflowVersionStreamSubject))
+      .pipe(takeUntil(this.componentDestroy))
       .subscribe({
         next: (response: TypeAnnotationResponse) => {
           const choices = response.choices || [];

From 2d53a09ccf7a19970f678fa37245f11658287271 Mon Sep 17 00:00:00 2001
From: IamtherealBrian <2638932112@qq.com>
Date: Tue, 24 Sep 2024 18:13:26 -0400
Subject: [PATCH 05/12] fix fmt

---
 .../code-editor.component.ts                  |  4 +-
 .../ai-assistant/ai-assistant.service.ts      | 47 ++++++++++---------
 2 files changed, 25 insertions(+), 26 deletions(-)

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 8cbaf356c4d..3132f9af9f0 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
@@ -62,8 +62,7 @@ export class CodeEditorComponent implements AfterViewInit, SafeStyle, OnDestroy
   // For "Add All Type Annotation" to show the UI individually
   private userResponseSubject?: Subject<void>;
   private isMultipleVariables: boolean = false;
-  private componentDestroy = new Subject<void>;
-
+  private componentDestroy = new Subject<void>();
 
   private generateLanguageTitle(language: string): string {
     return `${language.charAt(0).toUpperCase()}${language.slice(1)} UDF`;
@@ -330,7 +329,6 @@ export class CodeEditorComponent implements AfterViewInit, SafeStyle, OnDestroy
     }
   }
 
-
   private handleTypeAnnotation(
     code: string,
     range: monaco.Range,
diff --git a/core/gui/src/app/workspace/service/ai-assistant/ai-assistant.service.ts b/core/gui/src/app/workspace/service/ai-assistant/ai-assistant.service.ts
index 313fcd60f3c..b8950799cba 100644
--- a/core/gui/src/app/workspace/service/ai-assistant/ai-assistant.service.ts
+++ b/core/gui/src/app/workspace/service/ai-assistant/ai-assistant.service.ts
@@ -39,7 +39,6 @@ interface UnannotatedArgumentResponse {
   };
 }
 
-
 // Define AI model type
 export const AI_ASSISTANT_API_BASE_URL = `${AppSettings.getApiEndpoint()}/aiassistant`;
 export const AI_MODEL = {
@@ -97,27 +96,29 @@ export class AIAssistantService {
   public locateUnannotated(selectedCode: string, startLine: number): Observable<UnannotatedArgument[]> {
     const requestBody = { selectedCode, startLine };
 
-    return this.http.post<UnannotatedArgumentResponse>(`${AI_ASSISTANT_API_BASE_URL}/annotate-argument`, requestBody).pipe(
-      map(response => {
-        if (response) {
-          return response.underlying.result.value.map(
-            (item: UnannotatedArgumentItem): UnannotatedArgument => ({
-              name: item.underlying.name.value,
-              startLine: item.underlying.startLine.value,
-              startColumn: item.underlying.startColumn.value,
-              endLine: item.underlying.endLine.value,
-              endColumn: item.underlying.endColumn.value,
-            })
-          );
-        } else {
-          console.error("Unexpected response format:", response);
-          return [];
-        }
-      }),
-      catchError((error: unknown) => {
-        console.error("Request to backend failed:", error);
-        return of([]);
-      })
-    );
+    return this.http
+      .post<UnannotatedArgumentResponse>(`${AI_ASSISTANT_API_BASE_URL}/annotate-argument`, requestBody)
+      .pipe(
+        map(response => {
+          if (response) {
+            return response.underlying.result.value.map(
+              (item: UnannotatedArgumentItem): UnannotatedArgument => ({
+                name: item.underlying.name.value,
+                startLine: item.underlying.startLine.value,
+                startColumn: item.underlying.startColumn.value,
+                endLine: item.underlying.endLine.value,
+                endColumn: item.underlying.endColumn.value,
+              })
+            );
+          } else {
+            console.error("Unexpected response format:", response);
+            return [];
+          }
+        }),
+        catchError((error: unknown) => {
+          console.error("Request to backend failed:", error);
+          return of([]);
+        })
+      );
   }
 }

From de1e6860f89ea73e2ebe38ec10221aef6980bb7f Mon Sep 17 00:00:00 2001
From: IamtherealBrian <2638932112@qq.com>
Date: Wed, 25 Sep 2024 01:51:11 -0400
Subject: [PATCH 06/12] fix fmt

---
 .../aiassistant/AiAssistantResource.scala     | 29 ++++++++++++++-----
 1 file changed, 22 insertions(+), 7 deletions(-)

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
index 613a57729b7..8a390b8cf8f 100644
--- 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
@@ -111,19 +111,34 @@ class AIAssistantResource {
   @Path("/annotate-argument")
   @Consumes(Array(MediaType.APPLICATION_JSON))
   def locateUnannotated(request: LocateUnannotatedRequest, @Auth user: SessionUser): Response = {
+    // Encoding the code to transmit multi-line code as a single command-line argument
     val encodedCode = Base64.getEncoder.encodeToString(request.selectedCode.getBytes("UTF-8"))
     val pythonScriptPath =
-      Paths.get("src", "main", "scala", "edu", "uci", "ics", "texera", "web", "resource", "aiassistant", "type_annotation_visitor.py").toString
+      Paths
+        .get(
+          "src",
+          "main",
+          "scala",
+          "edu",
+          "uci",
+          "ics",
+          "texera",
+          "web",
+          "resource",
+          "aiassistant",
+          "type_annotation_visitor.py"
+        )
+        .toString
 
     try {
       val command = s"""python $pythonScriptPath "$encodedCode" ${request.startLine}"""
       val result = command.!!
-    val parsedResult = objectMapper.readValue(result, classOf[List[List[Any]]]).map {
-    case List(name: String, startLine: Int, startColumn: Int, endLine: Int, endColumn: Int) =>
-      UnannotatedArgument(name, startLine, startColumn, endLine, endColumn)
-    case _ =>
-      throw new RuntimeException("Unexpected format in Python script result")
-    }
+      val parsedResult = objectMapper.readValue(result, classOf[List[List[Any]]]).map {
+        case List(name: String, startLine: Int, startColumn: Int, endLine: Int, endColumn: Int) =>
+          UnannotatedArgument(name, startLine, startColumn, endLine, endColumn)
+        case _ =>
+          throw new RuntimeException("Unexpected format in Python script result")
+      }
       Response.ok(Json.obj("result" -> Json.toJson(parsedResult))).build()
     } catch {
       case e: Exception =>

From cad30301f1304c08c6d0cfc442a93077b69ddd7b Mon Sep 17 00:00:00 2001
From: IamtherealBrian <2638932112@qq.com>
Date: Fri, 27 Sep 2024 18:21:11 -0400
Subject: [PATCH 07/12] fix type_annotation_visitor.py and add test case

---
 .../test_type_annotation_visitor.py           | 138 ++++++++++++++++++
 .../aiassistant/type_annotation_visitor.py    |  73 +++++++--
 2 files changed, 196 insertions(+), 15 deletions(-)
 create mode 100644 core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/test_type_annotation_visitor.py

diff --git a/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/test_type_annotation_visitor.py b/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/test_type_annotation_visitor.py
new file mode 100644
index 00000000000..78034a0a075
--- /dev/null
+++ b/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/test_type_annotation_visitor.py
@@ -0,0 +1,138 @@
+import pytest
+from type_annotation_visitor import find_untyped_variables
+
+class TestFunctionsAndMethods:
+
+    # 1. Global Functions Test
+    @pytest.fixture
+    def global_functions_code(self):
+        return """def global_function(a, b=2, /, c=3, *, d, e=5, **kwargs):
+    pass
+
+def global_function_no_return(a, b):
+    return a + b
+
+def global_function_with_return(a: int, b: int) -> int:
+    return a + b
+"""
+
+    def test_global_functions(self, global_functions_code):
+        expected_result = [
+            ["c", 1, 32, 1, 33],
+            ["a", 1, 21, 1, 22],
+            ["b", 1, 24, 1, 25],
+            ["d", 1, 40, 1, 41],
+            ["e", 1, 43, 1, 44],
+            ["kwargs", 1, 50, 1, 56],
+            ["a", 4, 31, 4, 32],
+            ["b", 4, 34, 4, 35]
+
+        ]
+        untyped_vars = find_untyped_variables(global_functions_code, 1)
+        assert untyped_vars == expected_result
+
+    # 2. Class Methods and Static Methods Test
+    @pytest.fixture
+    def class_methods_code(self):
+        return """class MyClass:
+    def instance_method_no_annotation(self, x, y):
+        pass
+
+    @staticmethod
+    def static_method(a, b, /, c=3, *, d, **kwargs):
+        pass
+
+    @staticmethod
+    def static_method_with_annotation(a: int, b: int, /, *, c: int = 5) -> int:
+        return a + b + c
+
+    @classmethod
+    def class_method(cls, value, /, *, option=True):
+        pass
+
+    @classmethod
+    def class_method_with_annotation(cls, value: str, /, *, flag: bool = False) -> str:
+        return value.upper()
+"""
+
+    def test_class_methods(self, class_methods_code):
+        expected_result = [
+            ["x", 2, 45, 2, 46],
+            ["y", 2, 48, 2, 49],
+            ["c", 6, 32, 6, 33],
+            ["a", 6, 23, 6, 24],
+            ["b", 6, 26, 6, 27],
+            ["d", 6, 40, 6, 41],
+            ["kwargs", 6, 45, 6, 51],
+            ["value", 14, 27, 14, 32],
+            ["option", 14, 40, 14, 46]
+        ]
+        untyped_vars = find_untyped_variables(class_methods_code, 1)
+        assert untyped_vars == expected_result
+
+    # 3. Lambda Functions Test
+    @pytest.fixture
+    def lambda_code(self):
+        return """lambda_function = lambda x, y, /, z=0, *, w=1: x + y + z + w
+lambda_function_with_annotation = lambda x: x * 2
+"""
+
+    def test_lambda_functions(self, lambda_code):
+        untyped_vars = find_untyped_variables(lambda_code, 1)
+        assert untyped_vars == []
+
+    # 4. Comprehensive Functions Test
+    @pytest.fixture
+    def comprehensive_functions_code(self):
+        return """def default_args_function(a, b=2, /, c=3, *, d=4):
+    pass
+
+def args_kwargs_function(*args, **kwargs):
+    pass
+
+def function_with_return_annotation(a: int, b: int, /, *, c: int = 0) -> int:
+    return a + b + c
+
+def function_without_return_annotation(a, b):
+    return a + b
+"""
+
+    def test_comprehensive(self, comprehensive_functions_code):
+        expected_result = [
+            ["c", 1, 38, 1, 39],
+            ["a", 1, 27, 1, 28],
+            ["b", 1, 30, 1, 31],
+            ["d", 1, 46, 1, 47],
+            ["args", 4, 27, 4, 31],
+            ["kwargs", 4, 35, 4, 41],
+            ["a", 10, 40, 10, 41],
+            ["b", 10, 43, 10, 44]
+        ]
+        untyped_vars = find_untyped_variables(comprehensive_functions_code, 1)
+        assert untyped_vars == expected_result
+
+    # 5. Multi-line Functions Test
+    @pytest.fixture
+    def multi_line_function_code(self):
+        return """def multi_line_function(
+    a,
+    b: int = 10,
+    /,
+    c: str = "hello",
+    *,
+    d,
+    e=20,
+    **kwargs
+):
+    pass
+"""
+
+    def test_multi_lines_argument(self, multi_line_function_code):
+        expected_result = [
+            ["a", 2, 5, 2, 6],
+            ["d", 7, 5, 7, 6],
+            ["e", 8, 5, 8, 6],
+            ["kwargs", 9, 7, 9, 13]
+        ]
+        untyped_vars = find_untyped_variables(multi_line_function_code, 1)
+        assert untyped_vars == expected_result
diff --git a/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/type_annotation_visitor.py b/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/type_annotation_visitor.py
index 391f409e8bd..183e94ed2ed 100644
--- a/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/type_annotation_visitor.py
+++ b/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/type_annotation_visitor.py
@@ -3,37 +3,80 @@
 import sys
 import base64
 
+class ParentNodeVisitor(ast.NodeVisitor):
+    def __init__(self):
+        self.parent = None
+
+    def generic_visit(self, node):
+        node.parent = self.parent
+        previous_parent = self.parent
+        self.parent = node
+        super().generic_visit(node)
+        self.parent = previous_parent
+
 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':
+    def visit(self, node):
+        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+            self.process_function(node)
+        self.generic_visit(node)
+
+    def process_function(self, node):
+        # Boolean to determine if it's a global function or a method
+        is_method = isinstance(node.parent, ast.ClassDef)
+        # Boolean to determine if it's a static method
+        is_staticmethod = False
+        if is_method and hasattr(node, 'decorator_list'):
+            for decorator in node.decorator_list:
+                if isinstance(decorator, ast.Name) and decorator.id == 'staticmethod':
+                    is_staticmethod = True
+                elif isinstance(decorator, ast.Attribute) and decorator.attr == 'staticmethod':
+                    is_staticmethod = True
+        args = node.args
+
+        all_args = []
+        all_args.extend(args.args)
+        # Positional-only
+        all_args.extend(args.posonlyargs)
+        # Keyword-only
+        all_args.extend(args.kwonlyargs)
+        # *args
+        if args.vararg:
+            all_args.append(args.vararg)
+        # **kwargs
+        if args.kwarg:
+            all_args.append(args.kwarg)
+
+        start_index = 0
+        # Skip the "self" or "cls"
+        if is_method and not is_staticmethod:
+            start_index = 1
+        for i, arg in enumerate(all_args):
+            if i < start_index:
                 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)
+                self.add_untyped_arg(arg)
+
+    def add_untyped_arg(self, arg):
+        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])
 
 def find_untyped_variables(source_code, start_line):
     tree = ast.parse(source_code)
+    ParentNodeVisitor().visit(tree)
     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
+    print(json.dumps(untyped_variables))

From ecaeb732720c2b5bc416d859e54bf65feb91897f Mon Sep 17 00:00:00 2001
From: IamtherealBrian <2638932112@qq.com>
Date: Fri, 27 Sep 2024 18:59:55 -0400
Subject: [PATCH 08/12] fix MVVM

---
 .../code-editor.component.ts                   | 18 +++---------------
 1 file changed, 3 insertions(+), 15 deletions(-)

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 3132f9af9f0..037f23eceec 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
@@ -409,7 +409,6 @@ export class CodeEditorComponent implements AfterViewInit, SafeStyle, OnDestroy
     }
   }
 
-  // Add the type annotation into monaco editor
   private insertTypeAnnotations(
     editor: monaco.editor.IStandaloneCodeEditor,
     selection: monaco.Selection,
@@ -417,20 +416,9 @@ export class CodeEditorComponent implements AfterViewInit, SafeStyle, OnDestroy
   ) {
     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]);
+    const insertPosition = new monaco.Position(endLineNumber, endColumn);
+    const insertOffset = editor.getModel()?.getOffsetAt(insertPosition) || 0;
+    this.code?.insert(insertOffset, annotations);
   }
 
   private connectLanguageServer() {

From e816f7169926784d2cb86bd5ce269b889a294753 Mon Sep 17 00:00:00 2001
From: IamtherealBrian <2638932112@qq.com>
Date: Wed, 2 Oct 2024 13:40:25 -0400
Subject: [PATCH 09/12] fix part comments

---
 .../aiassistant/type_annotation_visitor.py    |  6 ++--
 .../ai-assistant/ai-assistant.service.ts      | 28 +++++++++++++------
 2 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/type_annotation_visitor.py b/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/type_annotation_visitor.py
index 183e94ed2ed..0ca26b0eb10 100644
--- a/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/type_annotation_visitor.py
+++ b/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/type_annotation_visitor.py
@@ -54,12 +54,11 @@ def process_function(self, node):
         # Skip the "self" or "cls"
         if is_method and not is_staticmethod:
             start_index = 1
-        for i, arg in enumerate(all_args):
-            if i < start_index:
-                continue
+        for i, arg in enumerate(all_args[start_index:]):
             if not arg.annotation:
                 self.add_untyped_arg(arg)
 
+
     def add_untyped_arg(self, arg):
         start_line = arg.lineno + self.start_line_offset - 1
         start_col = arg.col_offset + 1
@@ -77,6 +76,7 @@ def find_untyped_variables(source_code, start_line):
 if __name__ == "__main__":
     encoded_code = sys.argv[1]
     start_line = int(sys.argv[2])
+    # Encoding the code to transmit multi-line code as a single command-line argument before, so we need to decode it here
     source_code = base64.b64decode(encoded_code).decode('utf-8')
     untyped_variables = find_untyped_variables(source_code, start_line)
     print(json.dumps(untyped_variables))
diff --git a/core/gui/src/app/workspace/service/ai-assistant/ai-assistant.service.ts b/core/gui/src/app/workspace/service/ai-assistant/ai-assistant.service.ts
index b8950799cba..5d5130a14b9 100644
--- a/core/gui/src/app/workspace/service/ai-assistant/ai-assistant.service.ts
+++ b/core/gui/src/app/workspace/service/ai-assistant/ai-assistant.service.ts
@@ -13,13 +13,13 @@ export type TypeAnnotationResponse = {
   }>;
 };
 
-export interface UnannotatedArgument {
-  readonly name: string;
-  readonly startLine: number;
-  readonly startColumn: number;
-  readonly endLine: number;
-  readonly endColumn: number;
-}
+export interface UnannotatedArgument extends Readonly<{
+  name: string;
+  startLine: number;
+  startColumn: number;
+  endLine: number;
+  endColumn: number;
+}> {}
 
 interface UnannotatedArgumentItem {
   readonly underlying: {
@@ -101,6 +101,18 @@ export class AIAssistantService {
       .pipe(
         map(response => {
           if (response) {
+
+            const result = response.underlying.result.value.map(
+              (item: UnannotatedArgumentItem): UnannotatedArgument => ({
+                name: item.underlying.name.value,
+                startLine: item.underlying.startLine.value,
+                startColumn: item.underlying.startColumn.value,
+                endLine: item.underlying.endLine.value,
+                endColumn: item.underlying.endColumn.value,
+              })
+            );
+            console.log("Unannotated Arguments:", result);
+
             return response.underlying.result.value.map(
               (item: UnannotatedArgumentItem): UnannotatedArgument => ({
                 name: item.underlying.name.value,
@@ -117,7 +129,7 @@ export class AIAssistantService {
         }),
         catchError((error: unknown) => {
           console.error("Request to backend failed:", error);
-          return of([]);
+          throw new Error("Request to backend failed");
         })
       );
   }

From f5589a1d6a7d04969adf6079139fef582afd8fb1 Mon Sep 17 00:00:00 2001
From: IamtherealBrian <2638932112@qq.com>
Date: Wed, 2 Oct 2024 14:30:01 -0400
Subject: [PATCH 10/12] fix comments

---
 .../aiassistant/test_type_annotation_visitor.py    | 14 ++++++++------
 .../aiassistant/type_annotation_visitor.py         |  2 ++
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/test_type_annotation_visitor.py b/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/test_type_annotation_visitor.py
index 78034a0a075..10612e76f16 100644
--- a/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/test_type_annotation_visitor.py
+++ b/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/test_type_annotation_visitor.py
@@ -3,9 +3,9 @@
 
 class TestFunctionsAndMethods:
 
-    # 1. Global Functions Test
     @pytest.fixture
     def global_functions_code(self):
+        """This is the test for global function"""
         return """def global_function(a, b=2, /, c=3, *, d, e=5, **kwargs):
     pass
 
@@ -34,6 +34,7 @@ def test_global_functions(self, global_functions_code):
     # 2. Class Methods and Static Methods Test
     @pytest.fixture
     def class_methods_code(self):
+        """This is the test for class methods and static methods"""
         return """class MyClass:
     def instance_method_no_annotation(self, x, y):
         pass
@@ -70,20 +71,21 @@ def test_class_methods(self, class_methods_code):
         untyped_vars = find_untyped_variables(class_methods_code, 1)
         assert untyped_vars == expected_result
 
-    # 3. Lambda Functions Test
     @pytest.fixture
     def lambda_code(self):
+        """This is the test for lambda function"""
         return """lambda_function = lambda x, y, /, z=0, *, w=1: x + y + z + w
 lambda_function_with_annotation = lambda x: x * 2
 """
 
     def test_lambda_functions(self, lambda_code):
-        untyped_vars = find_untyped_variables(lambda_code, 1)
-        assert untyped_vars == []
+        with pytest.raises(ValueError) as exc_info:
+                find_untyped_variables(lambda_code, 1)
+        assert "Lambda functions do not support type annotation" in str(exc_info.value)
 
-    # 4. Comprehensive Functions Test
     @pytest.fixture
     def comprehensive_functions_code(self):
+        """This is the test for comprehensive function"""
         return """def default_args_function(a, b=2, /, c=3, *, d=4):
     pass
 
@@ -111,9 +113,9 @@ def test_comprehensive(self, comprehensive_functions_code):
         untyped_vars = find_untyped_variables(comprehensive_functions_code, 1)
         assert untyped_vars == expected_result
 
-    # 5. Multi-line Functions Test
     @pytest.fixture
     def multi_line_function_code(self):
+        """This is the test for multi-line function"""
         return """def multi_line_function(
     a,
     b: int = 10,
diff --git a/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/type_annotation_visitor.py b/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/type_annotation_visitor.py
index 0ca26b0eb10..39089960554 100644
--- a/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/type_annotation_visitor.py
+++ b/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/type_annotation_visitor.py
@@ -22,6 +22,8 @@ def __init__(self, start_line_offset=0):
     def visit(self, node):
         if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
             self.process_function(node)
+        elif isinstance(node, ast.Lambda):
+            raise ValueError("Lambda functions do not support type annotation")
         self.generic_visit(node)
 
     def process_function(self, node):

From a38ba567d873e505d49a87fac9772fd507a7064c Mon Sep 17 00:00:00 2001
From: IamtherealBrian <2638932112@qq.com>
Date: Wed, 2 Oct 2024 14:31:52 -0400
Subject: [PATCH 11/12] fix fmt

---
 .../service/ai-assistant/ai-assistant.service.ts | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/core/gui/src/app/workspace/service/ai-assistant/ai-assistant.service.ts b/core/gui/src/app/workspace/service/ai-assistant/ai-assistant.service.ts
index 5d5130a14b9..1c378896128 100644
--- a/core/gui/src/app/workspace/service/ai-assistant/ai-assistant.service.ts
+++ b/core/gui/src/app/workspace/service/ai-assistant/ai-assistant.service.ts
@@ -13,13 +13,14 @@ export type TypeAnnotationResponse = {
   }>;
 };
 
-export interface UnannotatedArgument extends Readonly<{
-  name: string;
-  startLine: number;
-  startColumn: number;
-  endLine: number;
-  endColumn: number;
-}> {}
+export interface UnannotatedArgument
+  extends Readonly<{
+    name: string;
+    startLine: number;
+    startColumn: number;
+    endLine: number;
+    endColumn: number;
+  }> {}
 
 interface UnannotatedArgumentItem {
   readonly underlying: {
@@ -101,7 +102,6 @@ export class AIAssistantService {
       .pipe(
         map(response => {
           if (response) {
-
             const result = response.underlying.result.value.map(
               (item: UnannotatedArgumentItem): UnannotatedArgument => ({
                 name: item.underlying.name.value,

From 6d034c901182023a951db3e327fbb08502745731 Mon Sep 17 00:00:00 2001
From: IamtherealBrian <2638932112@qq.com>
Date: Wed, 2 Oct 2024 14:35:31 -0400
Subject: [PATCH 12/12] fix comments

---
 .../web/resource/aiassistant/test_type_annotation_visitor.py     | 1 -
 1 file changed, 1 deletion(-)

diff --git a/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/test_type_annotation_visitor.py b/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/test_type_annotation_visitor.py
index 10612e76f16..4f6a4b0d3f6 100644
--- a/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/test_type_annotation_visitor.py
+++ b/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/test_type_annotation_visitor.py
@@ -31,7 +31,6 @@ def test_global_functions(self, global_functions_code):
         untyped_vars = find_untyped_variables(global_functions_code, 1)
         assert untyped_vars == expected_result
 
-    # 2. Class Methods and Static Methods Test
     @pytest.fixture
     def class_methods_code(self):
         """This is the test for class methods and static methods"""