Add AI-Assistant flag and features

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"
+}

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")
 }

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

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
+  }
+}

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)
+}

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))

core/gui/src/app/dashboard/service/user/ai-assistant/ai-assistant.service.ts

@@ -0,0 +1,78 @@
+import { Injectable } from "@angular/core";
+import { firstValueFrom } from "rxjs";
+import { AppSettings } from "../../../../common/app-setting";
+import { HttpClient } from "@angular/common/http";
+
+export const AI_ASSISTANT_API_BASE_URL = `${AppSettings.getApiEndpoint()}/aiassistant`;
+
+@Injectable({
+  providedIn: "root",
+})
+export class AiAssistantService {
+  constructor(private http: HttpClient) {}
+
+  public checkAiAssistantEnabled(): Promise<boolean> {
+    const apiUrl = `${AI_ASSISTANT_API_BASE_URL}/isenabled`;
+    return firstValueFrom(this.http.get<boolean>(apiUrl))
+      .then(response => (response !== undefined ? response : false))
+      .catch(() => false);
+  }
+
+  public getTypeAnnotations(code: string, lineNumber: number, allcode: string): Promise<string> {
+    const prompt = `
+            Your task is to analyze the given Python code and provide only the type annotation as stated in the instructions.
+            Instructions:
+            - The provided code will only be one of the 2 situations below:
+            - First situation: The input is not start with "def". If the provided code only contains variable, output the result in the format ":type".
+            - Second situation: The input is start with "def". If the provided code starts with "def" (a longer line than just a variable, indicative of a function or method), output the result in the format " -> type".
+            - The type should only be one word, such as "str", "int", etc.
+
+            Examples:
+            - First situation:
+                - Provided code is "name", then the output may be : str
+                - Provided code is "age", then the output may be : int
+                - Provided code is "data", then the output may be : Tuple[int, str]
+                - Provided code is "new_user", then the output may be : User
+                - A special case: provided code is "self" and the context is something like "def __init__(self, username :str , age :int)", if the user requires the type annotation for the first parameter "self", then you should generate nothing.
+            - Second situation: (actual output depends on the complete code content)
+                - Provided code is "process_data(data: List[Tuple[int, str]], config: Dict[str, Union[int, str]])", then the output may be -> Optional[str]
+                - Provided code is "def add(a: int, b: int)", then the output may be -> int
+
+            Counterexamples:
+            - Provided code is "def __init__(self, username: str, age: int)" and you generate the result:
+                The result is The provided code is "def __init__(self, username: str, age: int)", so it fits the second situation, which means the result should be in " -> type" format. However, the __init__ method in Python doesn't return anything or in other words, it implicitly returns None. Hence the correct type hint would be: -> None.
+                I don't want this result! The correct result you should generate is -> None for this counter case.
+
+            Details:
+            - Provided code: ${code}
+            - Line number of the provided code in the complete code context: ${lineNumber}
+            - Complete code context: ${allcode}
+
+            Important: (you must follow!!)
+            - For the first situation: you must return strictly according to the format ": type", without adding any extra characters. No need for an explanation, just the result : type is enough!
+            - For the second situation: you return strictly according to the format " -> type", without adding any extra characters. No need for an explanation, just the result -> type is enough!
+        `;
+    return firstValueFrom(this.http.post<any>(`${AI_ASSISTANT_API_BASE_URL}/getresult`, { prompt }))
+      .then(response => {
+        console.log("Received response from backend:", response);
+        const result = response.choices[0].message.content.trim();
+        return result;
+      })
+      .catch(error => {
+        console.error("Request to backend failed:", error);
+        return "";
+      });
+  }
+
+  public locateUnannotated(selectedCode: string, startLine: number) {
+    return firstValueFrom(this.http.post<any>(`${AI_ASSISTANT_API_BASE_URL}/getArgument`, { selectedCode, startLine }))
+      .then(response => {
+        console.log("Received response from backend:", response);
+        return response.result;
+      })
+      .catch(error => {
+        console.error("Request to backend failed:", error);
+        return [];
+      });
+  }
+}

core/gui/src/app/workspace/component/code-editor-dialog/code-editor.component.html

@@ -29,3 +29,33 @@
     <div [innerHTML]="this.getCoeditorCursorStyles(user)"></div>
   </ng-container>
 </div>
+
+<div
+  *ngIf="showAnnotationSuggestion"
+  class="annotation-suggestion">
+  <p>Do you agree with the type annotation suggestion？</p>
+  <pre>Adding annotation for code: {{ currentCode }}</pre>
+  <p>Given suggestion: <strong>{{ currentSuggestion }}</strong></p>
+  <button
+    class="accept-button"
+    (click)="acceptCurrentAnnotation()">
+    Accept
+  </button>
+  <button
+    class="decline-button"
+    (click)="rejectCurrentAnnotation()">
+    Decline
+  </button>
+</div>
+
+<div
+  id="noAnnotationNeeded"
+  class="noAnnotationNeeded"
+  style="display: none">
+  <h5 class="title">No Type Annotations Needed</h5>
+  <button
+    id="close-button-for-none"
+    class="close-button-for-none">
+    Ok
+  </button>
+</div>