refactor

apps/app/src/features/openai/server/routes/edit/schema.ts → apps/app/src/features/openai/interfaces/editor-assistant/llm-response-schemas.ts

@@ -5,11 +5,11 @@ import { z } from 'zod';
 // -----------------------------------------------------------------------------
 
 // Schema definitions
-export const EditorAssistantMessageSchema = z.object({
+export const LlmEditorAssistantMessageSchema = z.object({
   message: z.string().describe('A friendly message explaining what changes were made or suggested'),
 });
 
-export const EditorAssistantDiffSchema = z
+export const LlmEditorAssistantDiffSchema = z
   .object({
     insert: z.string().describe('The text that should insert the content in the current position'),
   })
@@ -23,3 +23,7 @@ export const EditorAssistantDiffSchema = z
       retain: z.number().int().describe('The number of characters that should be retained in the current position'),
     }),
   );
+
+// Type definitions
+export type LlmEditorAssistantMessage = z.infer<typeof LlmEditorAssistantMessageSchema>;
+export type LlmEditorAssistantDiff = z.infer<typeof LlmEditorAssistantDiffSchema>;

apps/app/src/features/openai/interfaces/editor-assistant/sse-schemas.ts

@@ -0,0 +1,30 @@
+import { z } from 'zod';
+
+import { LlmEditorAssistantDiffSchema } from './llm-response-schemas';
+
+// -----------------------------------------------------------------------------
+// Type definitions
+// -----------------------------------------------------------------------------
+
+// Schema definitions
+export const SseMessageSchema = z.object({
+  appendedMessage: z.string().describe('The message that should be appended to the chat window'),
+});
+
+export const SseDetectedDiffSchema = z
+  .object({
+    diff: LlmEditorAssistantDiffSchema,
+  });
+
+export const SseFinalizedSchema = z
+  .object({
+    finalized: z.object({
+      message: z.string().describe('The final message that should be displayed in the chat window'),
+      replacements: z.array(LlmEditorAssistantDiffSchema),
+    }),
+  });
+
+// Type definitions
+export type SseMessage = z.infer<typeof SseMessageSchema>;
+export type SseDetectedDiff = z.infer<typeof SseDetectedDiffSchema>;
+export type SseFinalized = z.infer<typeof SseFinalizedSchema>;

apps/app/src/features/openai/server/routes/edit/index.ts

@@ -8,23 +8,24 @@ import type { MessageDelta } from 'openai/resources/beta/threads/messages.mjs';
 import { z } from 'zod';
 
 // Necessary imports
-import { getOrCreateEditorAssistant } from '~/features/openai/server/services/assistant';
 import type Crowi from '~/server/crowi';
 import { accessTokenParser } from '~/server/middlewares/access-token-parser';
 import { apiV3FormValidator } from '~/server/middlewares/apiv3-form-validator';
 import type { ApiV3Response } from '~/server/routes/apiv3/interfaces/apiv3-response';
 import loggerFactory from '~/utils/logger';
 
+import { LlmEditorAssistantDiffSchema, LlmEditorAssistantMessageSchema } from '../../../interfaces/editor-assistant/llm-response-schemas';
+import type { SseDetectedDiff, SseFinalized, SseMessage } from '../../../interfaces/editor-assistant/sse-schemas';
 import { MessageErrorCode } from '../../../interfaces/message-error';
+import { getOrCreateEditorAssistant } from '../../services/assistant';
 import { openaiClient } from '../../services/client';
+import { LlmResponseStreamProcessor } from '../../services/editor-assistant';
 import { getStreamErrorCode } from '../../services/getStreamErrorCode';
 import { getOpenaiService } from '../../services/openai';
 import { replaceAnnotationWithPageLink } from '../../services/replace-annotation-with-page-link';
 import { certifyAiService } from '../middlewares/certify-ai-service';
 import { SseHelper } from '../utils/sse-helper';
 
-import { EditorStreamProcessor } from './editor-stream-processor';
-import { EditorAssistantDiffSchema, EditorAssistantMessageSchema } from './schema';
 
 const logger = loggerFactory('growi:routes:apiv3:openai:message');
 
@@ -32,8 +33,8 @@ const logger = loggerFactory('growi:routes:apiv3:openai:message');
 // Type definitions
 // -----------------------------------------------------------------------------
 
-const EditorAssistantResponseSchema = z.object({
-  contents: z.array(z.union([EditorAssistantMessageSchema, EditorAssistantDiffSchema])),
+const LlmEditorAssistantResponseSchema = z.object({
+  contents: z.array(z.union([LlmEditorAssistantMessageSchema, LlmEditorAssistantDiffSchema])),
 }).describe('The response format for the editor assistant');
 
 
@@ -95,7 +96,17 @@ export const postMessageToEditHandlersFactory: PostMessageHandlersFactory = (cro
 
       // Initialize SSE helper and stream processor
       const sseHelper = new SseHelper(res);
-      const streamProcessor = new EditorStreamProcessor(sseHelper);
+      const streamProcessor = new LlmResponseStreamProcessor({
+        messageCallback: (appendedMessage) => {
+          sseHelper.writeData<SseMessage>({ appendedMessage });
+        },
+        diffDetectedCallback: (detected) => {
+          sseHelper.writeData<SseDetectedDiff>({ diff: detected });
+        },
+        dataFinalizedCallback: (message, replacements) => {
+          sseHelper.writeData<SseFinalized>({ finalized: { message: message ?? '', replacements } });
+        },
+      });
 
       try {
         // Set response headers
@@ -146,7 +157,7 @@ export const postMessageToEditHandlersFactory: PostMessageHandlersFactory = (cro
               content: `Current markdown content:\n\`\`\`markdown\n${markdown}\n\`\`\`\n\nUser request: ${userMessage}`,
             },
           ],
-          response_format: zodResponseFormat(EditorAssistantResponseSchema, 'editor_assistant_response'),
+          response_format: zodResponseFormat(LlmEditorAssistantResponseSchema, 'editor_assistant_response'),
         });
 
         // Message delta handler
@@ -161,10 +172,11 @@ export const postMessageToEditHandlersFactory: PostMessageHandlersFactory = (cro
           // Process text
           if (content?.type === 'text' && content.text?.value) {
             const chunk = content.text.value;
-            rawBuffer += chunk;
 
             // Process data with JSON processor
-            streamProcessor.process(rawBuffer);
+            streamProcessor.process(rawBuffer, chunk);
+
+            rawBuffer += chunk;
 
             // Also send original delta
             sseHelper.writeData(delta);

apps/app/src/features/openai/server/routes/utils/sse-helper.ts

@@ -9,7 +9,7 @@ export interface ISseHelper {
   /**
    * Send data in SSE format
    */
-  writeData(data: unknown): void;
+  writeData<T extends object>(data: T): void;
 
   /**
    * Send error in SSE format
@@ -35,7 +35,7 @@ export class SseHelper implements ISseHelper {
   /**
    * Send data in SSE format
    */
-  writeData(data: unknown): void {
+  writeData<T extends object>(data: T): void {
     this.res.write(`data: ${JSON.stringify(data)}\n\n`);
   }
 

apps/app/src/features/openai/server/services/editor-assistant/index.ts

@@ -0,0 +1 @@
+export * from './llm-response-stream-processor';

apps/app/src/features/openai/server/routes/edit/editor-stream-processor.ts → apps/app/src/features/openai/server/services/editor-assistant/llm-response-stream-processor.ts

@@ -3,42 +3,43 @@ import type { z } from 'zod';
 
 import loggerFactory from '~/utils/logger';
 
-import type { SseHelper } from '../utils/sse-helper';
-
-import type { EditorAssistantMessageSchema } from './schema';
-import { EditorAssistantDiffSchema } from './schema';
+import {
+  type LlmEditorAssistantMessage,
+  LlmEditorAssistantDiffSchema, type LlmEditorAssistantDiff,
+} from '../../../interfaces/editor-assistant/llm-response-schemas';
 
 const logger = loggerFactory('growi:routes:apiv3:openai:edit:editor-stream-processor');
 
-// Type definitions
-type EditorAssistantMessage = z.infer<typeof EditorAssistantMessageSchema>;
-type EditorAssistantDiff = z.infer<typeof EditorAssistantDiffSchema>;
-
 /**
  * Type guard: Check if item is a message type
  */
-const isMessageItem = (item: unknown): item is EditorAssistantMessage => {
+const isMessageItem = (item: unknown): item is LlmEditorAssistantMessage => {
   return typeof item === 'object' && item !== null && 'message' in item;
 };
 
 /**
  * Type guard: Check if item is a diff type
  */
-const isDiffItem = (item: unknown): item is EditorAssistantDiff => {
+const isDiffItem = (item: unknown): item is LlmEditorAssistantDiff => {
   return typeof item === 'object' && item !== null
     && ('insert' in item || 'delete' in item || 'retain' in item);
 };
 
+type Options = {
+  messageCallback?: (appendedMessage: string) => void,
+  diffDetectedCallback?: (detected: LlmEditorAssistantDiff) => void,
+  dataFinalizedCallback?: (message: string | null, replacements: LlmEditorAssistantDiff[]) => void,
+}
 /**
- * Editor Stream Processor
+ * AI response stream processor for Editor Assisntant
  * Extracts messages and diffs from JSON stream for editor
  */
-export class EditorStreamProcessor {
+export class LlmResponseStreamProcessor {
 
   // Final response data
   private message: string | null = null;
 
-  private replacements: EditorAssistantDiff[] = [];
+  private replacements: LlmEditorAssistantDiff[] = [];
 
   // Index of the last element in previous content
   private lastContentIndex = -1;
@@ -49,15 +50,20 @@ export class EditorStreamProcessor {
   // Set of sent diff keys
   private sentDiffKeys = new Set<string>();
 
-  constructor(private sseHelper: SseHelper) {
-    this.sseHelper = sseHelper;
+  constructor(
+      private options?: Options,
+  ) {
+    this.options = options;
   }
 
   /**
    * Process JSON data
-   * @param jsonString JSON string
+   * @param prevJsonString Previous JSON string
+   * @param chunk New chunk of JSON string
    */
-  process(jsonString: string): void {
+  process(prevJsonString: string, chunk: string): void {
+    const jsonString = prevJsonString + chunk;
+
     try {
       const repairedJson = jsonrepair(jsonString);
       const parsedJson = JSON.parse(repairedJson);
@@ -90,11 +96,11 @@ export class EditorStreamProcessor {
           const item = contents[i];
           if (!isDiffItem(item)) continue;
 
-          const validDiff = EditorAssistantDiffSchema.safeParse(item);
+          const validDiff = LlmEditorAssistantDiffSchema.safeParse(item);
           if (!validDiff.success) continue;
 
           const diff = validDiff.data;
-          const key = this.getDiffKey(diff);
+          const key = this.getDiffKey(diff, i);
 
           // Check if this diff has already been sent
           if (this.sentDiffKeys.has(key)) continue;
@@ -113,14 +119,16 @@ export class EditorStreamProcessor {
         this.lastContentIndex = currentContentIndex;
 
         // Send notifications
-        if (messageUpdated) {
+        // TODO: Invoke callback with only appended strings
+        if (messageUpdated && this.message != null) {
           // Notify immediately if message is updated
-          this.notifyClient();
+          this.options?.messageCallback?.(this.message);
         }
-        else if (diffUpdated && processedDiffIndex > this.lastSentDiffIndex) {
+
+        if (diffUpdated && processedDiffIndex > this.lastSentDiffIndex) {
           // For diffs, only notify if a new index diff is confirmed
           this.lastSentDiffIndex = processedDiffIndex;
-          this.notifyClient();
+          this.options?.diffDetectedCallback?.(this.replacements[this.replacements.length - 1]);
         }
       }
     }
@@ -133,25 +141,13 @@ export class EditorStreamProcessor {
   /**
    * Generate unique key for a diff
    */
-  private getDiffKey(diff: EditorAssistantDiff): string {
-    if ('insert' in diff) return `insert-${diff.insert}`;
-    if ('delete' in diff) return `delete-${diff.delete}`;
-    if ('retain' in diff) return `retain-${diff.retain}`;
+  private getDiffKey(diff: LlmEditorAssistantDiff, index: number): string {
+    if ('insert' in diff) return `insert-${index}`;
+    if ('delete' in diff) return `delete-${index}`;
+    if ('retain' in diff) return `retain-${index}`;
     return '';
   }
 
-  /**
-   * Notify the client
-   */
-  private notifyClient(): void {
-    this.sseHelper.writeData({
-      editorResponse: {
-        message: this.message || '',
-        replacements: this.replacements,
-      },
-    });
-  }
-
   /**
    * Send final result
    */
@@ -165,14 +161,14 @@ export class EditorStreamProcessor {
         const contents = parsedJson.contents;
 
         // Add any unsent diffs
-        for (const item of contents) {
+        for (const [index, item] of contents) {
           if (!isDiffItem(item)) continue;
 
-          const validDiff = EditorAssistantDiffSchema.safeParse(item);
+          const validDiff = LlmEditorAssistantDiffSchema.safeParse(item);
           if (!validDiff.success) continue;
 
           const diff = validDiff.data;
-          const key = this.getDiffKey(diff);
+          const key = this.getDiffKey(diff, index);
 
           // Add any diffs that haven't been sent yet
           if (!this.sentDiffKeys.has(key)) {
@@ -182,26 +178,14 @@ export class EditorStreamProcessor {
         }
       }
 
-      // Final notification (with isDone flag)
-      this.sseHelper.writeData({
-        editorResponse: {
-          message: this.message || '',
-          replacements: this.replacements,
-        },
-        isDone: true,
-      });
+      // Final notification
+      this.options?.dataFinalizedCallback?.(this.message, this.replacements);
     }
     catch (e) {
       logger.debug('Failed to parse final JSON response:', e);
 
       // Send final notification even on error
-      this.sseHelper.writeData({
-        editorResponse: {
-          message: this.message || '',
-          replacements: this.replacements,
-        },
-        isDone: true,
-      });
+      this.options?.dataFinalizedCallback?.(this.message, this.replacements);
     }
   }