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 28802934329..2d7db7dbc0f 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 @@ -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]) } /** 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 4d4c3002d79..f784165dc65 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,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 + } } diff --git a/core/gui/src/app/app.module.ts b/core/gui/src/app/app.module.ts index 76bb1e7c991..060b4733ac3 100644 --- a/core/gui/src/app/app.module.ts +++ b/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, 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 deleted file mode 100644 index 66534152827..00000000000 --- a/core/gui/src/app/dashboard/service/user/ai-assistant/ai-assistant.service.ts +++ /dev/null @@ -1,30 +0,0 @@ -import { Injectable } from "@angular/core"; -import { firstValueFrom } from "rxjs"; -import { AppSettings } from "../../../../common/app-setting"; -import { HttpClient } from "@angular/common/http"; - -export const AI_ASSISTANT_API_BASE_URL = `${AppSettings.getApiEndpoint()}/aiassistant`; - -@Injectable({ - providedIn: "root", -}) -export class AiAssistantService { - constructor(private http: HttpClient) {} - - public checkAiAssistantEnabled(): Promise { - const apiUrl = `${AI_ASSISTANT_API_BASE_URL}/isenabled`; - return firstValueFrom(this.http.get(apiUrl, { 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" - ); - return isEnabled; - }) - .catch(() => { - return "NoAiAssistant"; - }); - } -} diff --git a/core/gui/src/app/workspace/component/code-editor-dialog/annotation-suggestion.component.html b/core/gui/src/app/workspace/component/code-editor-dialog/annotation-suggestion.component.html new file mode 100644 index 00000000000..9346ed85df6 --- /dev/null +++ b/core/gui/src/app/workspace/component/code-editor-dialog/annotation-suggestion.component.html @@ -0,0 +1,18 @@ +
+

Do you agree with the type annotation suggestion?

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

Given suggestion: {{ suggestion }}

+ + +
diff --git a/core/gui/src/app/workspace/component/code-editor-dialog/annotation-suggestion.component.scss b/core/gui/src/app/workspace/component/code-editor-dialog/annotation-suggestion.component.scss new file mode 100644 index 00000000000..b986983e1d2 --- /dev/null +++ b/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; +} diff --git a/core/gui/src/app/workspace/component/code-editor-dialog/annotation-suggestion.component.ts b/core/gui/src/app/workspace/component/code-editor-dialog/annotation-suggestion.component.ts new file mode 100644 index 00000000000..282f1b9578c --- /dev/null +++ b/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(); + @Output() decline = new EventEmitter(); + + onAccept() { + this.accept.emit(); + } + + onDecline() { + this.decline.emit(); + } +} diff --git a/core/gui/src/app/workspace/component/code-editor-dialog/code-editor.component.html b/core/gui/src/app/workspace/component/code-editor-dialog/code-editor.component.html index e94d617541e..63f615a31ac 100644 --- a/core/gui/src/app/workspace/component/code-editor-dialog/code-editor.component.html +++ b/core/gui/src/app/workspace/component/code-editor-dialog/code-editor.component.html @@ -29,3 +29,13 @@
+ + + 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..c772ad7786c 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,8 @@ 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, TypeAnnotationResponse } from "../../service/ai-assistant/ai-assistant.service"; +import { AnnotationSuggestionComponent } from "./annotation-suggestion.component"; /** * CodeEditorComponent is the content of the dialogue invoked by CodeareaCustomTemplateComponent. @@ -35,6 +37,7 @@ import { FormControl } from "@angular/forms"; export class CodeEditorComponent implements AfterViewInit, SafeStyle, OnDestroy { @ViewChild("editor", { static: true }) editorElement!: ElementRef; @ViewChild("container", { static: true }) containerElement!: ElementRef; + @ViewChild(AnnotationSuggestionComponent) annotationSuggestion!: AnnotationSuggestionComponent; private code?: YText; private editor?: any; private languageServerSocket?: WebSocket; @@ -46,6 +49,17 @@ export class CodeEditorComponent implements AfterViewInit, SafeStyle, OnDestroy public language: string = ""; public languageTitle: string = ""; + // Boolean to determine whether the suggestion UI should be shown + public showAnnotationSuggestion: boolean = false; + // The code selected by the user + public currentCode: string = ""; + // The result returned by the backend AI assistant + public currentSuggestion: string = ""; + // The range selected by the user + public currentRange: monaco.Range | undefined; + public suggestionTop: number = 0; + public suggestionLeft: number = 0; + private generateLanguageTitle(language: string): string { return `${language.charAt(0).toUpperCase()}${language.slice(1)} UDF`; } @@ -63,7 +77,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; @@ -169,11 +184,144 @@ export class CodeEditorComponent implements AfterViewInit, SafeStyle, OnDestroy ); } this.editor = editor; + + // Check if the AI provider is "openai" + this.aiAssistantService + .checkAIAssistantEnabled() + .pipe(untilDestroyed(this)) + .subscribe({ + next: (isEnabled: string) => { + if (isEnabled === "OpenAI") { + // "Add Type Annotation" Button + editor.addAction({ + id: "type-annotation-action", + label: "Add Type Annotation", + contextMenuGroupId: "1_modification", + contextMenuOrder: 1.0, + run: ed => { + // User selected code (including range and content) + const selection = ed.getSelection(); + const model = ed.getModel(); + if (!model || !selection) { + return; + } + // All the code in Python UDF + const allcode = model.getValue(); + // Content of user selected code + const code = model.getValueInRange(selection); + // Start line of the selected code + const lineNumber = selection.startLineNumber; + this.handleTypeAnnotation( + code, + selection, + ed as monaco.editor.IStandaloneCodeEditor, + lineNumber, + allcode + ); + }, + }); + } + }, + }); if (this.language == "python") { this.connectLanguageServer(); } } + private handleTypeAnnotation( + code: string, + range: monaco.Range, + editor: monaco.editor.IStandaloneCodeEditor, + lineNumber: number, + allcode: string + ): void { + this.aiAssistantService + .getTypeAnnotations(code, lineNumber, allcode) + .pipe(takeUntil(this.workflowVersionStreamSubject)) + .subscribe({ + next: (response: TypeAnnotationResponse) => { + const choices = response.choices || []; + if (choices.length > 0 && choices[0].message && choices[0].message.content) { + this.currentSuggestion = choices[0].message.content.trim(); + this.currentCode = code; + this.currentRange = range; + + const position = editor.getScrolledVisiblePosition(range.getStartPosition()); + if (position) { + this.suggestionTop = position.top + 100; + this.suggestionLeft = position.left + 100; + } + + this.showAnnotationSuggestion = true; + + if (this.annotationSuggestion) { + this.annotationSuggestion.code = this.currentCode; + this.annotationSuggestion.suggestion = this.currentSuggestion; + this.annotationSuggestion.top = this.suggestionTop; + this.annotationSuggestion.left = this.suggestionLeft; + } + } else { + console.error("Error: OpenAI response does not contain valid message content", response); + } + }, + error: (error: unknown) => { + console.error("Error fetching type annotations:", error); + }, + }); + } + + // Called when the user clicks the "accept" button + public acceptCurrentAnnotation(): void { + // Avoid accidental calls + if (!this.showAnnotationSuggestion || !this.currentRange || !this.currentSuggestion) { + return; + } + + if (this.currentRange && this.currentSuggestion) { + const selection = new monaco.Selection( + this.currentRange.startLineNumber, + this.currentRange.startColumn, + this.currentRange.endLineNumber, + this.currentRange.endColumn + ); + this.insertTypeAnnotations(this.editor, selection, this.currentSuggestion); + } + // close the UI after adding the annotation + this.showAnnotationSuggestion = false; + } + + // Called when the user clicks the "decline" button + public rejectCurrentAnnotation(): void { + // Do nothing except for closing the UI + this.showAnnotationSuggestion = false; + this.currentCode = ""; + this.currentSuggestion = ""; + } + + // Add the type annotation into monaco editor + private insertTypeAnnotations( + editor: monaco.editor.IStandaloneCodeEditor, + selection: monaco.Selection, + annotations: string + ) { + const endLineNumber = selection.endLineNumber; + const endColumn = selection.endColumn; + const range = new monaco.Range( + // Insert the content to the end of the selected code + endLineNumber, + endColumn, + endLineNumber, + endColumn + ); + const text = `${annotations}`; + const op = { + range: range, + text: text, + forceMoveMarkers: true, + }; + editor.executeEdits("add annotation", [op]); + } + private connectLanguageServer() { if (this.languageServerSocket === undefined) { this.languageServerSocket = new WebSocket(getWebsocketUrl("/python-language-server", "3000")); 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 new file mode 100644 index 00000000000..0463b6af309 --- /dev/null +++ b/core/gui/src/app/workspace/service/ai-assistant/ai-assistant.service.ts @@ -0,0 +1,69 @@ +import { Injectable } from "@angular/core"; +import { AppSettings } from "../../../common/app-setting"; +import { HttpClient, HttpHeaders } from "@angular/common/http"; +import { Observable, of } from "rxjs"; +import { map, catchError } from "rxjs/operators"; + +// The type annotation return from the LLM +export type TypeAnnotationResponse = { + choices: ReadonlyArray<{ + message: Readonly<{ + content: string; + }>; + }>; +}; + +// Define AI model type +export const AI_ASSISTANT_API_BASE_URL = `${AppSettings.getApiEndpoint()}/aiassistant`; +export const AI_MODEL = { + OpenAI: "OpenAI", + NoAiAssistant: "NoAiAssistant", +} as const; +export type AI_MODEL = (typeof AI_MODEL)[keyof typeof AI_MODEL]; + +@Injectable({ + providedIn: "root", +}) +export class AIAssistantService { + constructor(private http: HttpClient) {} + + /** + * Checks if AI Assistant is enabled and returns the AI model in use. + * + * @returns {Observable} - An Observable that emits the type of AI model in use ("OpenAI" or "NoAiAssistant"). + */ + // To get the backend AI flag to check if the user want to use the AI feature + // valid returns: ["OpenAI", "NoAiAssistant"] + public checkAIAssistantEnabled(): Observable { + const apiUrl = `${AI_ASSISTANT_API_BASE_URL}/isenabled`; + return this.http.get(apiUrl, { responseType: "text" }).pipe( + map(response => { + const isEnabled: AI_MODEL = response === "OpenAI" ? "OpenAI" : "NoAiAssistant"; + console.log( + isEnabled === "OpenAI" + ? "AI Assistant successfully started" + : "No AI Assistant or OpenAI authentication key error" + ); + return isEnabled; + }), + catchError(() => { + return of("NoAiAssistant" as AI_MODEL); + }) + ); + } + + /** + * Sends a request to the backend to get type annotation suggestions from LLM for the provided code. + * + * @param {string} code - The selected code for which the user wants type annotation suggestions. + * @param {number} lineNumber - The line number where the selected code locates. + * @param {string} allcode - The entire code of the UDF (User Defined Function) to provide context for the AI assistant. + * + * @returns {Observable} - An Observable that emits the type annotation suggestions + * returned by the LLM. + */ + public getTypeAnnotations(code: string, lineNumber: number, allcode: string): Observable { + const requestBody = { code, lineNumber, allcode }; + return this.http.post(`${AI_ASSISTANT_API_BASE_URL}/annotationresult`, requestBody, {}); + } +}