diff --git a/ai/services/service.py b/ai/services/service.py index 7a4a97fbb..cf4da94da 100644 --- a/ai/services/service.py +++ b/ai/services/service.py @@ -1,3 +1,4 @@ +import json import os import re import secrets @@ -5,7 +6,7 @@ from os import getenv from typing import NamedTuple -from flask import Flask, Response, request +from flask import Flask, Response, request, stream_with_context from openai import OpenAI app = Flask(__name__) @@ -53,7 +54,7 @@ def save_interaction_endpoint() -> Response | dict[str, str]: @app.route("/adjust-mesop-app", methods=["POST"]) -def adjust_mesop_app_endpoint() -> Response | dict[str, str]: +def adjust_mesop_app_endpoint(): data = request.json assert data is not None code = data.get("code") @@ -62,14 +63,23 @@ def adjust_mesop_app_endpoint() -> Response | dict[str, str]: if not code or not prompt: return Response("Both 'code' and 'prompt' are required", status=400) - try: - diff = adjust_mesop_app(code, prompt) + def generate(): + stream = adjust_mesop_app(code, prompt) + diff = "" + for chunk in stream: + if chunk.choices[0].delta.content: + diff += chunk.choices[0].delta.content + yield f"data: {json.dumps({'type': 'progress', 'data': chunk.choices[0].delta.content})}\n\n" + result = apply_patch(code, diff) if result.has_error: raise Exception(result.result) - return {"code": result.result, "diff": diff} - except Exception as e: - return Response(f"Error: {e!s}", status=500) + + yield f"data: {json.dumps({'type': 'end', 'code': result.result, 'diff': diff})}\n\n" + + return Response( + stream_with_context(generate()), content_type="text/event-stream" + ) class ApplyPatchResult(NamedTuple): @@ -98,18 +108,12 @@ def apply_patch(original_code: str, patch: str) -> ApplyPatchResult: return ApplyPatchResult(False, patched_code) -def adjust_mesop_app(code: str, msg: str) -> str: +def adjust_mesop_app(code: str, msg: str): model = "ft:gpt-4o-mini-2024-07-18:personal::9yoxJtKf" client = OpenAI( api_key=getenv("OPENAI_API_KEY"), ) - return adjust_mesop_app_openai_client(code, msg, client, model=model) - - -def adjust_mesop_app_openai_client( - code: str, msg: str, client: OpenAI, model: str -) -> str: - completion = client.chat.completions.create( + return client.chat.completions.create( model=model, max_tokens=10_000, messages=[ @@ -124,11 +128,8 @@ def adjust_mesop_app_openai_client( ), }, ], + stream=True, ) - print("[INFO] LLM output:", completion.choices[0].message.content) - llm_output = completion.choices[0].message.content - assert llm_output is not None - return llm_output def generate_folder_name(prompt: str) -> str: diff --git a/mesop/server/server.py b/mesop/server/server.py index 64761dc22..8ce472fb5 100644 --- a/mesop/server/server.py +++ b/mesop/server/server.py @@ -5,7 +5,7 @@ import secrets import time import urllib.parse as urlparse -from typing import Generator, Sequence +from typing import Any, Generator, Sequence from urllib import request as urllib_request from urllib.error import URLError @@ -331,7 +331,7 @@ def save_interaction() -> Response | dict[str, str]: ) @flask_app.route("/__editor__/page-generate", methods=["POST"]) - def page_generate() -> Response | dict[str, str]: + def page_generate(): check_editor_access() try: @@ -357,38 +357,30 @@ def page_generate() -> Response | dict[str, str]: source_code = file.read() print(f"Source code of module {module.__name__}:") - try: - req = urllib_request.Request( + def generate(): + for event in sse_request( AI_SERVICE_BASE_URL + "/adjust-mesop-app", - data=json.dumps({"prompt": prompt, "code": source_code}).encode( - "utf-8" - ), - headers={"Content-Type": "application/json"}, - ) - with urllib_request.urlopen(req) as response: - if response.status == 200: - response_data = json.loads(response.read().decode("utf-8")) - generated_code = response_data["code"] - diff = response_data["diff"] + {"prompt": prompt, "code": source_code}, + ): + if event.get("type") == "end": + sse_data = { + "type": "end", + "prompt": prompt, + "path": path, + "beforeCode": source_code, + "afterCode": event["code"], + "diff": event["diff"], + "message": "Prompt processed successfully", + } + yield f"data: {json.dumps(sse_data)}\n\n" + break + elif event.get("type") == "progress": + sse_data = {"data": event["data"], "type": "progress"} + yield f"data: {json.dumps(sse_data)}\n\n" else: - print(f"Error from AI service: {response.read().decode('utf-8')}") - return Response( - f"Error from AI service: {response.read().decode('utf-8')}", - status=500, - ) - except URLError as e: - return Response( - f"Error making request to AI service: {e!s}", status=500 - ) + raise Exception(f"Unknown event type: {event}") - return { - "prompt": prompt, - "path": path, - "beforeCode": source_code, - "afterCode": generated_code, - "diff": diff, - "message": "Prompt processed successfully", - } + return Response(generate(), content_type="text/event-stream") @flask_app.route("/__hot-reload__") def hot_reload() -> Response: @@ -466,3 +458,30 @@ def is_same_site(url1: str | None, url2: str | None): return p1.hostname == p2.hostname except ValueError: return False + + +SSE_DATA_PREFIX = "data: " + + +def sse_request( + url: str, data: dict[str, Any] +) -> Generator[dict[str, Any], None, None]: + """ + Make an SSE request and yield JSON parsed events. + """ + headers = { + "Content-Type": "application/json", + "Accept": "text/event-stream", + } + encoded_data = json.dumps(data).encode("utf-8") + req = urllib_request.Request( + url, data=encoded_data, headers=headers, method="POST" + ) + + with urllib_request.urlopen(req) as response: + for line in response: + if line.strip(): + decoded_line = line.decode("utf-8").strip() + if decoded_line.startswith(SSE_DATA_PREFIX): + event_data = json.loads(decoded_line[len(SSE_DATA_PREFIX) :]) + yield event_data diff --git a/mesop/web/src/editor_toolbar/code_mirror_component.ts b/mesop/web/src/editor_toolbar/code_mirror_component.ts index 3f674f167..29db31383 100644 --- a/mesop/web/src/editor_toolbar/code_mirror_component.ts +++ b/mesop/web/src/editor_toolbar/code_mirror_component.ts @@ -1,6 +1,7 @@ import {EditorView, basicSetup} from 'codemirror'; import {python} from '@codemirror/lang-python'; -import {MergeView} from '@codemirror/merge'; +import {unifiedMergeView} from '@codemirror/merge'; +import {EditorState} from '@codemirror/state'; import { Component, @@ -11,47 +12,77 @@ import { } from '@angular/core'; @Component({ - selector: 'mesop-code-mirror', + selector: 'mesop-code-mirror-diff', template: '', standalone: true, }) -export class CodeMirrorComponent { +export class CodeMirrorDiffComponent { @Input() beforeCode!: string; @Input() afterCode!: string; @Output() codeChange = new EventEmitter(); + view: EditorView | null = null; constructor(private elementRef: ElementRef) {} ngOnChanges() { - while (this.elementRef.nativeElement.firstChild) { - this.elementRef.nativeElement.firstChild.remove(); + if (this.view) { + this.view.destroy(); } this.renderEditor(); } renderEditor() { - const mergeView = new MergeView({ - a: { - doc: this.beforeCode, + this.view = new EditorView({ + state: EditorState.create({ + doc: this.afterCode, extensions: [ basicSetup, python(), EditorView.editable.of(false), EditorView.lineWrapping, + unifiedMergeView({ + original: this.beforeCode, + highlightChanges: true, + mergeControls: false, + gutter: true, + collapseUnchanged: {margin: 2}, + }), ], - }, - b: { - doc: this.afterCode, + }), + parent: this.elementRef.nativeElement, + }); + } +} + +@Component({ + selector: 'mesop-code-mirror-raw', + template: '', + standalone: true, +}) +export class CodeMirrorRawComponent { + @Input() code!: string | null; + private view: EditorView | null = null; + + constructor(private elementRef: ElementRef) {} + + ngOnChanges() { + if (this.view) { + this.view.destroy(); + } + this.renderEditor(); + } + + renderEditor() { + this.view = new EditorView({ + state: EditorState.create({ + doc: this.code ?? '', extensions: [ basicSetup, python(), EditorView.editable.of(false), EditorView.lineWrapping, ], - }, + }), parent: this.elementRef.nativeElement, - highlightChanges: true, - collapseUnchanged: {margin: 2}, - gutter: true, }); } } diff --git a/mesop/web/src/editor_toolbar/editor_history_dialog.ng.html b/mesop/web/src/editor_toolbar/editor_history_dialog.ng.html index 18238fb58..a78535541 100644 --- a/mesop/web/src/editor_toolbar/editor_history_dialog.ng.html +++ b/mesop/web/src/editor_toolbar/editor_history_dialog.ng.html @@ -22,11 +22,10 @@

Mesop Editor History

> save - - + /> } diff --git a/mesop/web/src/editor_toolbar/editor_response_dialog.ng.html b/mesop/web/src/editor_toolbar/editor_response_dialog.ng.html index b717fe4cf..ce0429990 100644 --- a/mesop/web/src/editor_toolbar/editor_response_dialog.ng.html +++ b/mesop/web/src/editor_toolbar/editor_response_dialog.ng.html @@ -1,9 +1,9 @@

Mesop Editor Response

- + /> diff --git a/mesop/web/src/editor_toolbar/editor_send_prompt_progress_dialog.ng.html b/mesop/web/src/editor_toolbar/editor_send_prompt_progress_dialog.ng.html new file mode 100644 index 000000000..b01965b12 --- /dev/null +++ b/mesop/web/src/editor_toolbar/editor_send_prompt_progress_dialog.ng.html @@ -0,0 +1,7 @@ +

Mesop Editor Progress

+ + + + + + diff --git a/mesop/web/src/editor_toolbar/editor_toolbar.ts b/mesop/web/src/editor_toolbar/editor_toolbar.ts index 18d84685a..c363cf4fa 100644 --- a/mesop/web/src/editor_toolbar/editor_toolbar.ts +++ b/mesop/web/src/editor_toolbar/editor_toolbar.ts @@ -13,7 +13,10 @@ import { MatDialog, MatDialogModule, } from '@angular/material/dialog'; -import {CodeMirrorComponent} from './code_mirror_component'; +import { + CodeMirrorDiffComponent, + CodeMirrorRawComponent, +} from './code_mirror_component'; import {interval, Subscription} from 'rxjs'; import {MatSnackBar, MatSnackBarModule} from '@angular/material/snack-bar'; import {CommonModule} from '@angular/common'; @@ -138,7 +141,16 @@ export class EditorToolbar implements OnInit { this.startTimer(startTime); try { - const response = await this.editorToolbarService.sendPrompt(prompt); + const responsePromise = this.editorToolbarService.sendPrompt(prompt); + const progressDialogRef = this.dialog.open( + EditorSendPromptProgressDialog, + { + width: '90%', + }, + ); + const response = await responsePromise; + progressDialogRef.close(); + this.autocompleteTrigger.closePanel(); const dialogRef = this.dialog.open(EditorPromptResponseDialog, { data: {response: response, responseTime: this.responseTime}, width: '90%', @@ -250,10 +262,28 @@ export class EditorToolbar implements OnInit { } } +@Component({ + templateUrl: 'editor_send_prompt_progress_dialog.ng.html', + standalone: true, + imports: [ + MatDialogModule, + MatButtonModule, + CommonModule, + CodeMirrorRawComponent, + ], +}) +class EditorSendPromptProgressDialog { + constructor(private editorToolbarService: EditorToolbarService) {} + + get progress$() { + return this.editorToolbarService.generationProgress$; + } +} + @Component({ templateUrl: 'editor_response_dialog.ng.html', standalone: true, - imports: [MatDialogModule, MatButtonModule, CodeMirrorComponent], + imports: [MatDialogModule, MatButtonModule, CodeMirrorDiffComponent], }) class EditorPromptResponseDialog { constructor( @@ -268,7 +298,7 @@ class EditorPromptResponseDialog { imports: [ MatDialogModule, MatButtonModule, - CodeMirrorComponent, + CodeMirrorDiffComponent, MatIconModule, MatSnackBarModule, MatTooltipModule, diff --git a/mesop/web/src/editor_toolbar/editor_toolbar_service.ts b/mesop/web/src/editor_toolbar/editor_toolbar_service.ts index f275f7279..e846412d1 100644 --- a/mesop/web/src/editor_toolbar/editor_toolbar_service.ts +++ b/mesop/web/src/editor_toolbar/editor_toolbar_service.ts @@ -1,4 +1,6 @@ -import {Injectable} from '@angular/core'; +import {Injectable, NgZone} from '@angular/core'; +import {SSE} from '../utils/sse'; +import {BehaviorSubject, Observable} from 'rxjs'; export interface PromptInteraction extends PromptResponse { readonly prompt: string; @@ -11,11 +13,28 @@ export interface PromptResponse { readonly diff: string; } +interface GenerateEndMessage extends PromptResponse { + readonly type: 'end'; +} + +interface GenerateProgressMessage { + readonly type: 'progress'; + readonly data: string; +} + +type GenerateData = GenerateEndMessage | GenerateProgressMessage; + @Injectable({ providedIn: 'root', }) export class EditorToolbarService { history: PromptInteraction[] = []; + eventSource: SSE | undefined; + private readonly generationProgressSubject = new BehaviorSubject(''); + readonly generationProgress$: Observable = + this.generationProgressSubject.asObservable(); + + constructor(private readonly ngZone: NgZone) {} getHistory(): readonly PromptInteraction[] { return this.history; @@ -23,23 +42,51 @@ export class EditorToolbarService { async sendPrompt(prompt: string): Promise { console.debug('sendPrompt', prompt); + // Clear the progress subject + this.generationProgressSubject.next(''); const path = window.location.pathname; - const response = await fetch('/__editor__/page-generate', { - method: 'POST', - headers: { - 'Content-Type': 'application/json', - }, - body: JSON.stringify({prompt, path}), - }); - await handleError(response); - const json = (await response.json()) as PromptResponse; - // Insert at the top of the history so we display the most recent interactions first. - this.history.unshift({ - path, - prompt, - ...json, + return new Promise((resolve, reject) => { + this.eventSource = new SSE('/__editor__/page-generate', { + payload: JSON.stringify({prompt, path}), + headers: { + 'Content-Type': 'application/json', + }, + }); + this.eventSource.addEventListener('message', (e) => { + // Looks like Angular has a bug where it's not intercepting EventSource onmessage. + this.ngZone.run(() => { + try { + const data = (e as any).data; + const obj = JSON.parse(data) as GenerateData; + if (!obj.type) { + reject(new Error('Invalid event source message')); + return; + } + if (obj.type === 'end') { + this.eventSource!.close(); + this.eventSource = undefined; + const {beforeCode, afterCode, diff} = obj; + this.history.unshift({ + path, + prompt, + beforeCode, + afterCode, + diff, + }); + resolve({beforeCode, afterCode, diff}); + } + if (obj.type === 'progress') { + this.generationProgressSubject.next( + this.generationProgressSubject.getValue() + obj.data, + ); + } + } catch (e) { + console.error('sendPrompt eventSource error', e); + reject(e); + } + }); + }); }); - return json; } async commit(code: string) { diff --git a/package.json b/package.json index 2cc045565..760923c92 100644 --- a/package.json +++ b/package.json @@ -59,7 +59,7 @@ "@angular/material-experimental": "^18.0.0", "@angular/platform-browser": "^18.0.0", "@codemirror/lang-python": "^6.1.6", - "@codemirror/merge": "^6.6.7", + "@codemirror/merge": "^6.7.0", "@types/google.maps": "^3.54.10", "@types/youtube": "^0.0.46", "codemirror": "^6.0.1", diff --git a/yarn.lock b/yarn.lock index 380b67bc3..f01739f01 100644 --- a/yarn.lock +++ b/yarn.lock @@ -2133,10 +2133,10 @@ "@codemirror/view" "^6.0.0" crelt "^1.0.5" -"@codemirror/merge@^6.6.7": - version "6.6.7" - resolved "https://registry.yarnpkg.com/@codemirror/merge/-/merge-6.6.7.tgz#deb1d2778c92b1bd6aa7205a3a17a8ead5c357f1" - integrity sha512-fgZHAuLuxIQi1U/oeszzJHAGlQfkGC3Rmd9/Lxs4yO9GUC798h9640aiPWTuAyY3+H2XmlzQcg5wfG9mObKqRQ== +"@codemirror/merge@^6.7.0": + version "6.7.0" + resolved "https://registry.yarnpkg.com/@codemirror/merge/-/merge-6.7.0.tgz#429d73d8aa64e06c2d3c5a56d9aa6a86cc2d3cb3" + integrity sha512-hVGhYZMBKLnb0Q8NXZ06uR1oFOuiMLNs/igZov5PpOqCaxQLGQA5y2zRojn1G6SyBQ0/nYM3aFeJb5kr4xlZag== dependencies: "@codemirror/language" "^6.0.0" "@codemirror/state" "^6.0.0"