Stream generate response to UI

ai/services/service.py

@@ -1,11 +1,12 @@
+import json
 import os
 import re
 import secrets
 import urllib.parse
 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:

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

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<string>();
+  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,
     });
   }
 }

mesop/web/src/editor_toolbar/editor_history_dialog.ng.html

@@ -22,11 +22,10 @@ <h3 mat-dialog-title>Mesop Editor History</h3>
       >
         <mat-icon>save</mat-icon>
       </button>
-      <mesop-code-mirror
+      <mesop-code-mirror-diff
         [beforeCode]="history[selectedInteraction].beforeCode"
         [afterCode]="history[selectedInteraction].afterCode"
-      >
-      </mesop-code-mirror>
+      />
       }
     </div>
   </div>

mesop/web/src/editor_toolbar/editor_response_dialog.ng.html

@@ -1,9 +1,9 @@
 <h3 mat-dialog-title>Mesop Editor Response</h3>
 <mat-dialog-content class="mat-typography">
-  <mesop-code-mirror
+  <mesop-code-mirror-diff
     [beforeCode]="data.response.beforeCode"
     [afterCode]="data.response.afterCode"
-  ></mesop-code-mirror>
+  />
 </mat-dialog-content>
 <mat-dialog-actions>
   <button mat-button [mat-dialog-close]="true">OK</button>

mesop/web/src/editor_toolbar/editor_send_prompt_progress_dialog.ng.html

@@ -0,0 +1,7 @@
+<h3 mat-dialog-title>Mesop Editor Progress</h3>
+<mat-dialog-content class="progress-code">
+  <mesop-code-mirror-raw [code]="progress$ | async"></mesop-code-mirror-raw>
+</mat-dialog-content>
+<mat-dialog-actions>
+  <button mat-button [mat-dialog-close]="true">OK</button>
+</mat-dialog-actions>

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,