[2/3] Add AI Assistant Service - add the “Add Type Annotation" Python UDF action

core/amber/src/main/scala/edu/uci/ics/texera/web/TexeraWebApplication.scala

@@ -260,7 +260,7 @@ class TexeraWebApplication
     environment.jersey.register(classOf[AdminExecutionResource])
     environment.jersey.register(classOf[UserQuotaResource])
     environment.jersey.register(classOf[UserDiscussionResource])
-    environment.jersey.register(classOf[AiAssistantResource])
+    environment.jersey.register(classOf[AIAssistantResource])
   }
 
   /**

core/amber/src/main/scala/edu/uci/ics/texera/web/resource/aiassistant/AiAssistantResource.scala

@@ -1,13 +1,89 @@
 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 javax.ws.rs.Consumes
+import javax.ws.rs.core.MediaType
+import play.api.libs.json.Json
+import kong.unirest.Unirest
+
+case class AIAssistantRequest(code: String, lineNumber: Int, allcode: String)
 
 @Path("/aiassistant")
-class AiAssistantResource {
+class AIAssistantResource {
   final private lazy val isEnabled = AiAssistantManager.validAIAssistant
   @GET
   @RolesAllowed(Array("REGULAR", "ADMIN"))
   @Path("/isenabled")
-  def isAiAssistantEnable: String = isEnabled
+  def isAIAssistantEnable: String = isEnabled
+
+  /**
+    * To get the type annotation suggestion from OpenAI
+    */
+  @POST
+  @RolesAllowed(Array("REGULAR", "ADMIN"))
+  @Path("/annotationresult")
+  @Consumes(Array(MediaType.APPLICATION_JSON))
+  def getAnnotation(
+      request: AIAssistantRequest,
+      @Auth user: SessionUser
+  ): Response = {
+    val finalPrompt = generatePrompt(request.code, request.lineNumber, request.allcode)
+    val requestBodyJson = Json.obj(
+      "model" -> "gpt-4",
+      "messages" -> Json.arr(
+        Json.obj(
+          "role" -> "user",
+          "content" -> finalPrompt
+        )
+      ),
+      "max_tokens" -> 15
+    )
+
+    val response = Unirest
+      .post(s"${AiAssistantManager.sharedUrl}/chat/completions")
+      .header("Authorization", s"Bearer ${AiAssistantManager.accountKey}")
+      .header("Content-Type", "application/json")
+      .body(Json.stringify(requestBodyJson))
+      .asString()
+    if (response.getStatus >= 400) {
+      throw new RuntimeException(s"getAnnotation error: ${response.getStatus}: ${response.getBody}")
+    }
+    Response.status(response.getStatus).entity(response.getBody).build()
+  }
+
+  // Helper function to get the type annotation
+  def generatePrompt(code: String, lineNumber: Int, allcode: String): String = {
+    s"""
+       |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.
+       |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!
+     """.stripMargin
+  }
 }

core/gui/src/app/app.module.ts

@@ -41,6 +41,7 @@ import { UserQuotaComponent } from "./dashboard/component/user/user-quota/user-q
 import { UserIconComponent } from "./dashboard/component/user/user-icon/user-icon.component";
 import { UserAvatarComponent } from "./dashboard/component/user/user-avatar/user-avatar.component";
 import { CodeEditorComponent } from "./workspace/component/code-editor-dialog/code-editor.component";
+import { AnnotationSuggestionComponent } from "./workspace/component/code-editor-dialog/annotation-suggestion.component";
 import { CodeareaCustomTemplateComponent } from "./workspace/component/codearea-custom-template/codearea-custom-template.component";
 import { MiniMapComponent } from "./workspace/component/workflow-editor/mini-map/mini-map.component";
 import { MenuComponent } from "./workspace/component/menu/menu.component";
@@ -168,6 +169,7 @@ registerLocaleData(en);
     VisualizationFrameContentComponent,
     CodeareaCustomTemplateComponent,
     CodeEditorComponent,
+    AnnotationSuggestionComponent,
     TypeCastingDisplayComponent,
     ShareAccessComponent,
     WorkflowExecutionHistoryComponent,

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

@@ -0,0 +1,18 @@
+<div
+  class="annotation-suggestion"
+  [style.top.px]="top"
+  [style.left.px]="left">
+  <p>Do you agree with the type annotation suggestion?</p>
+  <pre>Adding annotation for code: {{ code }}</pre>
+  <p>Given suggestion: <strong>{{ suggestion }}</strong></p>
+  <button
+    class="accept-button"
+    (click)="onAccept()">
+    Accept
+  </button>
+  <button
+    class="decline-button"
+    (click)="onDecline()">
+    Decline
+  </button>
+</div>

core/gui/src/app/workspace/component/code-editor-dialog/annotation-suggestion.component.scss

@@ -0,0 +1,39 @@
+.annotation-suggestion {
+  position: absolute;
+  background: #222;
+  color: #fff;
+  padding: 20px;
+  border-radius: 8px;
+  box-shadow: 0px 0px 10px rgba(0, 0, 0, 0.5);
+  z-index: 1000;
+}
+
+.annotation-suggestion button {
+  margin-right: 10px;
+}
+
+.annotation-suggestion button.accept-button {
+  background-color: #28a745;
+  color: #000;
+  border: none;
+  padding: 10px 20px;
+  border-radius: 5px;
+  cursor: pointer;
+}
+
+.annotation-suggestion button.accept-button:hover {
+  background-color: #218838;
+}
+
+.annotation-suggestion button.decline-button {
+  background-color: #dc3545;
+  color: #000;
+  border: none;
+  padding: 10px 20px;
+  border-radius: 5px;
+  cursor: pointer;
+}
+
+.annotation-suggestion button.decline-button:hover {
+  background-color: #c82333;
+}

core/gui/src/app/workspace/component/code-editor-dialog/annotation-suggestion.component.ts

@@ -0,0 +1,23 @@
+import { Component, Input, Output, EventEmitter } from "@angular/core";
+
+@Component({
+  selector: "texera-annotation-suggestion",
+  templateUrl: "./annotation-suggestion.component.html",
+  styleUrls: ["./annotation-suggestion.component.scss"],
+})
+export class AnnotationSuggestionComponent {
+  @Input() code: string = "";
+  @Input() suggestion: string = "";
+  @Input() top: number = 0;
+  @Input() left: number = 0;
+  @Output() accept = new EventEmitter<void>();
+  @Output() decline = new EventEmitter<void>();
+
+  onAccept() {
+    this.accept.emit();
+  }
+
+  onDecline() {
+    this.decline.emit();
+  }
+}

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

@@ -29,3 +29,13 @@
     <div [innerHTML]="this.getCoeditorCursorStyles(user)"></div>
   </ng-container>
 </div>
+
+<texera-annotation-suggestion
+  *ngIf="showAnnotationSuggestion"
+  [code]="currentCode"
+  [suggestion]="currentSuggestion"
+  [top]="suggestionTop"
+  [left]="suggestionLeft"
+  (accept)="acceptCurrentAnnotation()"
+  (decline)="rejectCurrentAnnotation()">
+</texera-annotation-suggestion>