From d85670b97daabb1b9ca53df286ddc1c40aaee98f Mon Sep 17 00:00:00 2001
From: "Minchong(Brian) Wu" <2638932112@qq.com>
Date: Wed, 28 Aug 2024 23:48:26 -0700
Subject: [PATCH 01/14] all

---
 .../amber/src/main/resources/application.conf |   5 +
 .../ics/amber/engine/common/AmberConfig.scala |   2 +
 .../aiassistant/AiAssistantManager.scala      |  43 ++++++
 .../aiassistant/AiAssistantResource.scala     | 132 ++++++++++++++++++
 .../python_abstract_syntax_tree.py            |  39 ++++++
 5 files changed, 221 insertions(+)
 create mode 100644 core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/AiAssistantManager.scala
 create mode 100644 core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/AiAssistantResource.scala
 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/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/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..62a1ce38739
--- /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
+  }
+}
\ No newline at end of file
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..e2e0c8af895
--- /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)
+}
\ No newline at end of file
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

From 0243d16f25ad1e7cc1553fb319fd1b8fdb76f402 Mon Sep 17 00:00:00 2001
From: "Minchong(Brian) Wu" <2638932112@qq.com>
Date: Thu, 29 Aug 2024 00:13:45 -0700
Subject: [PATCH 02/14] miss a file

---
 .../main/scala/edu/uci/ics/texera/web/TexeraWebApplication.scala | 1 +
 1 file changed, 1 insertion(+)

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])
   }
 
   /**

From b4fcb23ebdc0214bf3f59e02f9c205be1b5ee735 Mon Sep 17 00:00:00 2001
From: "Minchong(Brian) Wu" <2638932112@qq.com>
Date: Thu, 29 Aug 2024 00:31:10 -0700
Subject: [PATCH 03/14] fmt

---
 .../aiassistant/AiAssistantManager.scala       |  2 +-
 .../aiassistant/AiAssistantResource.scala      | 18 +++++++++---------
 2 files changed, 10 insertions(+), 10 deletions(-)

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
index 62a1ce38739..465adf46461 100644
--- 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
@@ -40,4 +40,4 @@ object AiAssistantManager {
     case _ =>
       false
   }
-}
\ No newline at end of file
+}
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 e2e0c8af895..e42055bf034 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
@@ -24,8 +24,8 @@ class AiAssistantResource {
   def isAiAssistantEnabled: Boolean = isEnabled
 
   /**
-   * To get the type annotation suggestion from OpenAI
-   */
+    * To get the type annotation suggestion from OpenAI
+    */
   @POST
   @RolesAllowed(Array("REGULAR", "ADMIN"))
   @Path("/getresult")
@@ -86,12 +86,12 @@ class AiAssistantResource {
         case Some(data: List[List[Any]]) =>
           val unannotatedArgs = data.map {
             case List(
-            name: String,
-            startLine: Double,
-            startColumn: Double,
-            endLine: Double,
-            endColumn: Double
-            ) =>
+                  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")
@@ -129,4 +129,4 @@ class AiAssistantResource {
 
 object AiAssistantResource {
   case class LocateUnannotatedRequest(selectedCode: String, startLine: Int)
-}
\ No newline at end of file
+}

From 817c9b49c611d7e7f4d2089ff47d5093a18588ae Mon Sep 17 00:00:00 2001
From: "Minchong(Brian) Wu" <2638932112@qq.com>
Date: Sat, 31 Aug 2024 15:44:43 -0400
Subject: [PATCH 04/14] fix review comments

---
 .../amber/src/main/resources/application.conf |   3 +-
 .../aiassistant/AiAssistantManager.scala      |   3 +-
 .../aiassistant/AiAssistantResource.scala     | 119 ------------------
 .../python_abstract_syntax_tree.py            |  39 ------
 4 files changed, 4 insertions(+), 160 deletions(-)
 delete 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/resources/application.conf b/core/amber/src/main/resources/application.conf
index 21f21a8af9a..32887a6b1fc 100644
--- a/core/amber/src/main/resources/application.conf
+++ b/core/amber/src/main/resources/application.conf
@@ -94,5 +94,6 @@ region-plan-generator {
 
 ai-assistant-server{
     assistant = "openai"
-    ai-service-key = "Put your OpenAI authentication key here"
+    # Put your OpenAI authentication key here
+    ai-service-key = ""
 }
\ No newline at end of file
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
index 465adf46461..d8c51c8e076 100644
--- 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
@@ -14,10 +14,11 @@ object AiAssistantManager {
       false
 
     case "openai" =>
+      val sharedUrl = "https://api.openai.com/v1"
       var isKeyValid: Boolean = false
       var connection: HttpURLConnection = null
       try {
-        val url = new URL("https://api.openai.com/v1/models")
+        val url = new URL(s"${sharedUrl}/models")
         connection = url.openConnection().asInstanceOf[HttpURLConnection]
         connection.setRequestMethod("GET")
         connection.setRequestProperty(
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 e42055bf034..50978ab2f37 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
@@ -1,132 +1,13 @@
 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
deleted file mode 100644
index 391f409e8bd..00000000000
--- a/core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/python_abstract_syntax_tree.py
+++ /dev/null
@@ -1,39 +0,0 @@
-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

From 09c9132b8db72b16c20890e0ce1bf7e101d9d9d7 Mon Sep 17 00:00:00 2001
From: "Minchong(Brian) Wu" <2638932112@qq.com>
Date: Sat, 31 Aug 2024 22:37:32 -0400
Subject: [PATCH 05/14] gui

---
 .../user/ai-assistant/ai-assistant.service.ts | 20 +++++++++++++++++++
 .../code-editor.component.ts                  |  6 ++++--
 2 files changed, 24 insertions(+), 2 deletions(-)
 create mode 100644 core/gui/src/app/dashboard/service/user/ai-assistant/ai-assistant.service.ts

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..224aa04d6e6
--- /dev/null
+++ b/core/gui/src/app/dashboard/service/user/ai-assistant/ai-assistant.service.ts
@@ -0,0 +1,20 @@
+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<boolean> {
+    const apiUrl = `${AI_ASSISTANT_API_BASE_URL}/isenabled`;
+    return firstValueFrom(this.http.get<boolean>(apiUrl))
+      .then(response => (response !== undefined ? response : false))
+      .catch(() => false);
+  }
+}
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..e42962c1b63 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.
@@ -63,7 +64,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 +154,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,

From 54789e9700dad7f7a085be23272729a34bde76c0 Mon Sep 17 00:00:00 2001
From: "Minchong(Brian) Wu" <2638932112@qq.com>
Date: Sat, 31 Aug 2024 22:43:29 -0400
Subject: [PATCH 06/14] gui

---
 .../component/code-editor-dialog/code-editor.component.ts       | 2 --
 1 file changed, 2 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 e42962c1b63..eb0b209738f 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,7 +17,6 @@ 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.
@@ -65,7 +64,6 @@ export class CodeEditorComponent implements AfterViewInit, SafeStyle, OnDestroy
     private workflowActionService: WorkflowActionService,
     private workflowVersionService: WorkflowVersionService,
     public coeditorPresenceService: CoeditorPresenceService,
-    private aiAssistantService: AiAssistantService
   ) {
     const currentOperatorId = this.workflowActionService.getJointGraphWrapper().getCurrentHighlightedOperatorIDs()[0];
     const operatorType = this.workflowActionService.getTexeraGraph().getOperator(currentOperatorId).operatorType;

From cfe231f26302406a58af5ca7d28362ca6c5cbc11 Mon Sep 17 00:00:00 2001
From: "Minchong(Brian) Wu" <2638932112@qq.com>
Date: Sat, 31 Aug 2024 22:45:10 -0400
Subject: [PATCH 07/14] gui

---
 .../component/code-editor-dialog/code-editor.component.ts       | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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 eb0b209738f..c21b7dadcd6 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
@@ -152,7 +152,7 @@ export class CodeEditorComponent implements AfterViewInit, SafeStyle, OnDestroy
    * Create a Monaco editor and connect it to MonacoBinding.
    * @private
    */
-  private async initMonaco() {
+  private initMonaco() {
     const editor = monaco.editor.create(this.editorElement.nativeElement, {
       language: this.language,
       fontSize: 11,

From 2a8ee30cfeebfdcae7898ab6c5f921a0aff17fbb Mon Sep 17 00:00:00 2001
From: "Minchong(Brian) Wu" <2638932112@qq.com>
Date: Sat, 31 Aug 2024 22:45:51 -0400
Subject: [PATCH 08/14] gui

---
 .../component/code-editor-dialog/code-editor.component.ts       | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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 c21b7dadcd6..d0ed9343d0e 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
@@ -63,7 +63,7 @@ export class CodeEditorComponent implements AfterViewInit, SafeStyle, OnDestroy
     private sanitizer: DomSanitizer,
     private workflowActionService: WorkflowActionService,
     private workflowVersionService: WorkflowVersionService,
-    public coeditorPresenceService: CoeditorPresenceService,
+    public coeditorPresenceService: CoeditorPresenceService
   ) {
     const currentOperatorId = this.workflowActionService.getJointGraphWrapper().getCurrentHighlightedOperatorIDs()[0];
     const operatorType = this.workflowActionService.getTexeraGraph().getOperator(currentOperatorId).operatorType;

From 232ada74d915ec6e251bf2589f110d62ee2ec66a Mon Sep 17 00:00:00 2001
From: "Minchong(Brian) Wu" <2638932112@qq.com>
Date: Sun, 1 Sep 2024 11:59:09 -0400
Subject: [PATCH 09/14] logger to frontend

---
 core/amber/src/main/resources/application.conf       |  2 +-
 .../resource/aiassistant/AiAssistantManager.scala    |  4 ----
 .../user/ai-assistant/ai-assistant.service.ts        | 12 ++++++++++--
 3 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/core/amber/src/main/resources/application.conf b/core/amber/src/main/resources/application.conf
index 32887a6b1fc..b0d9a943dad 100644
--- a/core/amber/src/main/resources/application.conf
+++ b/core/amber/src/main/resources/application.conf
@@ -93,7 +93,7 @@ region-plan-generator {
 }
 
 ai-assistant-server{
-    assistant = "openai"
+    assistant = "none"
     # Put your OpenAI authentication key here
     ai-service-key = ""
 }
\ No newline at end of file
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
index d8c51c8e076..5200a982f59 100644
--- 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
@@ -1,11 +1,8 @@
 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")
@@ -30,7 +27,6 @@ object AiAssistantManager {
       } catch {
         case e: Exception =>
           isKeyValid = false
-          logger.warning(s"Error validating OpenAI API key: ${e.getMessage}")
       } finally {
         if (connection != null) {
           connection.disconnect()
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
index 224aa04d6e6..ebe35bb1a16 100644
--- 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
@@ -14,7 +14,15 @@ export class AiAssistantService {
   public checkAiAssistantEnabled(): Promise<boolean> {
     const apiUrl = `${AI_ASSISTANT_API_BASE_URL}/isenabled`;
     return firstValueFrom(this.http.get<boolean>(apiUrl))
-      .then(response => (response !== undefined ? response : false))
-      .catch(() => false);
+      .then(response => {
+        const isEnabled = response !== undefined ? response : false;
+        console.log(isEnabled ? "AI Assistant successfully started" : "No AI Assistant or OpenAI authentication key error");
+        return isEnabled;
+      })
+      .catch(() => {
+        console.log("No AI Assistant or OpenAI authentication key error");
+        return false;
+      });
   }
+
 }

From e55837fa439671d47a23ed6ad333f6f442098ac8 Mon Sep 17 00:00:00 2001
From: "Minchong(Brian) Wu" <2638932112@qq.com>
Date: Sun, 1 Sep 2024 13:49:36 -0400
Subject: [PATCH 10/14] frontend fmt

---
 .../service/user/ai-assistant/ai-assistant.service.ts        | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

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
index ebe35bb1a16..3b45177d1b1 100644
--- 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
@@ -16,7 +16,9 @@ export class AiAssistantService {
     return firstValueFrom(this.http.get<boolean>(apiUrl))
       .then(response => {
         const isEnabled = response !== undefined ? response : false;
-        console.log(isEnabled ? "AI Assistant successfully started" : "No AI Assistant or OpenAI authentication key error");
+        console.log(
+          isEnabled ? "AI Assistant successfully started" : "No AI Assistant or OpenAI authentication key error"
+        );
         return isEnabled;
       })
       .catch(() => {
@@ -24,5 +26,4 @@ export class AiAssistantService {
         return false;
       });
   }
-
 }

From 038be99c02e51a2ddeb51cebd4c8b55f4328751e Mon Sep 17 00:00:00 2001
From: "Minchong(Brian) Wu" <2638932112@qq.com>
Date: Sun, 1 Sep 2024 19:52:11 -0400
Subject: [PATCH 11/14] fix comment

---
 .../amber/src/main/resources/application.conf |  4 +-
 .../aiassistant/AiAssistantManager.scala      | 51 +++++++++++--------
 2 files changed, 32 insertions(+), 23 deletions(-)

diff --git a/core/amber/src/main/resources/application.conf b/core/amber/src/main/resources/application.conf
index b0d9a943dad..460a60a9310 100644
--- a/core/amber/src/main/resources/application.conf
+++ b/core/amber/src/main/resources/application.conf
@@ -94,6 +94,8 @@ region-plan-generator {
 
 ai-assistant-server{
     assistant = "none"
-    # Put your OpenAI authentication key here
+    # Put your Ai Service authentication key here
     ai-service-key = ""
+    # Put your Ai service url here (If you are using OpenAI, then the url should be "https://api.openai.com/v1")
+    ai-service-url = ""
 }
\ No newline at end of file
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
index 5200a982f59..efb29f6d827 100644
--- 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
@@ -5,34 +5,41 @@ import java.net.{HttpURLConnection, URL}
 object AiAssistantManager {
   private val aiAssistantConfig = AmberConfig.aiAssistantConfig
   val assistantType: String = aiAssistantConfig.getString("assistant")
+  // The accountKey is the OpenAI authentication key used to authenticate API requests and obtain responses from the OpenAI service.
+
   val accountKey: String = aiAssistantConfig.getString("ai-service-key")
+  val sharedUrl: String = aiAssistantConfig.getString("ai-service-url")
+
+  private def initOpenAI(): Boolean = {
+    var isKeyValid: Boolean = false
+    var connection: HttpURLConnection = null
+    try {
+      val url = new URL(s"${sharedUrl}/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
+    } finally {
+      if (connection != null) {
+        connection.disconnect()
+      }
+    }
+    isKeyValid
+  }
+
   val validAIAssistant: Boolean = assistantType match {
     case "none" =>
       false
 
     case "openai" =>
-      val sharedUrl = "https://api.openai.com/v1"
-      var isKeyValid: Boolean = false
-      var connection: HttpURLConnection = null
-      try {
-        val url = new URL(s"${sharedUrl}/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
-      } finally {
-        if (connection != null) {
-          connection.disconnect()
-        }
-      }
-      isKeyValid
+      initOpenAI()
 
     case _ =>
       false

From 52f0eb840a90950ea2d7e69a114c953245668ece Mon Sep 17 00:00:00 2001
From: "Minchong(Brian) Wu" <2638932112@qq.com>
Date: Sun, 1 Sep 2024 20:35:46 -0400
Subject: [PATCH 12/14] fix comment

---
 .../aiassistant/AiAssistantManager.scala       | 18 ++++++++++--------
 .../aiassistant/AiAssistantResource.scala      |  2 +-
 .../user/ai-assistant/ai-assistant.service.ts  | 13 ++++++-------
 3 files changed, 17 insertions(+), 16 deletions(-)

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
index efb29f6d827..4ab4ced668e 100644
--- 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
@@ -10,8 +10,7 @@ object AiAssistantManager {
   val accountKey: String = aiAssistantConfig.getString("ai-service-key")
   val sharedUrl: String = aiAssistantConfig.getString("ai-service-url")
 
-  private def initOpenAI(): Boolean = {
-    var isKeyValid: Boolean = false
+  private def initOpenAI(): String = {
     var connection: HttpURLConnection = null
     try {
       val url = new URL(s"${sharedUrl}/models")
@@ -22,26 +21,29 @@ object AiAssistantManager {
         s"Bearer ${accountKey.trim.replaceAll("^\"|\"$", "")}"
       )
       val responseCode = connection.getResponseCode
-      isKeyValid = responseCode == 200
+      if (responseCode == 200) {
+        "OpenAI"
+      } else {
+        "NoAiAssistant"
+      }
     } catch {
       case e: Exception =>
-        isKeyValid = false
+        "NoAiAssistant"
     } finally {
       if (connection != null) {
         connection.disconnect()
       }
     }
-    isKeyValid
   }
 
-  val validAIAssistant: Boolean = assistantType match {
+  val validAIAssistant: String = assistantType match {
     case "none" =>
-      false
+      "NoAiAssistant"
 
     case "openai" =>
       initOpenAI()
 
     case _ =>
-      false
+      "NoAiAssistant"
   }
 }
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 50978ab2f37..4d4c3002d79 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,5 +9,5 @@ class AiAssistantResource {
   @GET
   @RolesAllowed(Array("REGULAR", "ADMIN"))
   @Path("/isenabled")
-  def isAiAssistantEnabled: Boolean = isEnabled
+  def isAiAssistantEnable: String = isEnabled
 }
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
index 3b45177d1b1..01b8fa3ecc8 100644
--- 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
@@ -11,19 +11,18 @@ export const AI_ASSISTANT_API_BASE_URL = `${AppSettings.getApiEndpoint()}/aiassi
 export class AiAssistantService {
   constructor(private http: HttpClient) {}
 
-  public checkAiAssistantEnabled(): Promise<boolean> {
+  public checkAiAssistantEnabled(): Promise<string> {
     const apiUrl = `${AI_ASSISTANT_API_BASE_URL}/isenabled`;
-    return firstValueFrom(this.http.get<boolean>(apiUrl))
+    return firstValueFrom(this.http.get(apiUrl, { responseType: "text"}))
       .then(response => {
-        const isEnabled = response !== undefined ? response : false;
+        const isEnabled = response !== undefined ? response : "NoAiAssistant";
         console.log(
-          isEnabled ? "AI Assistant successfully started" : "No AI Assistant or OpenAI authentication key error"
+          isEnabled === "OpenAI" ? "AI Assistant successfully started" : "No AI Assistant or OpenAI authentication key error"
         );
         return isEnabled;
       })
-      .catch(() => {
-        console.log("No AI Assistant or OpenAI authentication key error");
-        return false;
+      .catch(()=> {
+        return "NoAiAssistant";
       });
   }
 }

From 5769815efdd86dae57d43d11501ce42d713240f0 Mon Sep 17 00:00:00 2001
From: "Minchong(Brian) Wu" <2638932112@qq.com>
Date: Sun, 1 Sep 2024 20:48:53 -0400
Subject: [PATCH 13/14] fmt

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

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
index 01b8fa3ecc8..66534152827 100644
--- 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
@@ -13,15 +13,17 @@ export class AiAssistantService {
 
   public checkAiAssistantEnabled(): Promise<string> {
     const apiUrl = `${AI_ASSISTANT_API_BASE_URL}/isenabled`;
-    return firstValueFrom(this.http.get(apiUrl, { responseType: "text"}))
+    return firstValueFrom(this.http.get(apiUrl, { responseType: "text" }))
       .then(response => {
         const isEnabled = response !== undefined ? response : "NoAiAssistant";
         console.log(
-          isEnabled === "OpenAI" ? "AI Assistant successfully started" : "No AI Assistant or OpenAI authentication key error"
+          isEnabled === "OpenAI"
+            ? "AI Assistant successfully started"
+            : "No AI Assistant or OpenAI authentication key error"
         );
         return isEnabled;
       })
-      .catch(()=> {
+      .catch(() => {
         return "NoAiAssistant";
       });
   }

From f52d87f556a3ac41715d618ef4a058f3cdfc0930 Mon Sep 17 00:00:00 2001
From: "Minchong(Brian) Wu" <2638932112@qq.com>
Date: Mon, 2 Sep 2024 00:39:52 -0400
Subject: [PATCH 14/14] error handling

---
 .../scala/edu/uci/ics/amber/engine/common/AmberConfig.scala  | 5 ++++-
 .../texera/web/resource/aiassistant/AiAssistantManager.scala | 4 +++-
 2 files changed, 7 insertions(+), 2 deletions(-)

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 a5e3b8b4469..f5691b5d402 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
@@ -106,5 +106,8 @@ object AmberConfig {
   val jdbcConfig: Config = getConfSource.getConfig("jdbc")
 
   // Python language server configuration
-  val aiAssistantConfig: Config = getConfSource.getConfig("ai-assistant-server")
+  var aiAssistantConfig: Option[Config] = None
+  if (getConfSource.hasPath("ai-assistant-server")) {
+    aiAssistantConfig = Some(getConfSource.getConfig("ai-assistant-server"))
+  }
 }
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
index 4ab4ced668e..00e334a17aa 100644
--- 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
@@ -3,7 +3,9 @@ import edu.uci.ics.amber.engine.common.AmberConfig
 import java.net.{HttpURLConnection, URL}
 
 object AiAssistantManager {
-  private val aiAssistantConfig = AmberConfig.aiAssistantConfig
+  private val aiAssistantConfig = AmberConfig.aiAssistantConfig.getOrElse(
+    throw new Exception("ai-assistant-server configuration is missing in application.conf")
+  )
   val assistantType: String = aiAssistantConfig.getString("assistant")
   // The accountKey is the OpenAI authentication key used to authenticate API requests and obtain responses from the OpenAI service.