Explorar o código

refactor getPageBodyForContext

Yuki Takei hai 9 meses
pai
achega
fa97dd3de1

+ 74 - 23
apps/app/src/features/openai/client/services/editor-assistant/get-page-body-for-context.spec.ts

@@ -39,10 +39,20 @@ describe('getPageBodyForContext', () => {
 
       const result = getPageBodyForContext(mockEditor, 100, 200);
 
-      // Should default cursor position to 0 and take 200 chars after
-      const expectedContent = longContent.slice(0, 200);
-      expect(result).toBe(expectedContent);
-      expect(result).toHaveLength(200);
+      // Should default cursor position to 0
+      // Available before cursor: 0 (cursor at start)
+      // Shortfall before: 100 - 0 = 100
+      // Chars after: 200 + 100 = 300
+      // Expected: start=0, end=0+300=300
+      const expectedContent = longContent.slice(0, 300);
+      expect(result).toEqual({
+        content: expectedContent,
+        isPartial: true,
+        startIndex: 0,
+        endIndex: 300,
+        totalLength: 1000,
+      });
+      expect(result?.content).toHaveLength(300);
     });
   });
 
@@ -57,7 +67,11 @@ describe('getPageBodyForContext', () => {
 
       const result = getPageBodyForContext(mockEditor, 10, 10);
 
-      expect(result).toBe(shortText);
+      expect(result).toEqual({
+        content: shortText,
+        isPartial: false,
+        totalLength: 5,
+      });
       expect(mockEditor.getDocString).toHaveBeenCalled();
     });
 
@@ -70,7 +84,11 @@ describe('getPageBodyForContext', () => {
 
       const result = getPageBodyForContext(mockEditor, 50, 100); // total: 150
 
-      expect(result).toBe(exactLengthText);
+      expect(result).toEqual({
+        content: exactLengthText,
+        isPartial: false,
+        totalLength: 150,
+      });
       expect(mockEditor.getDocString).toHaveBeenCalled();
     });
 
@@ -83,7 +101,11 @@ describe('getPageBodyForContext', () => {
 
       const result = getPageBodyForContext(mockEditor, 50, 100); // total: 150
 
-      expect(result).toBe(shortText);
+      expect(result).toEqual({
+        content: shortText,
+        isPartial: false,
+        totalLength: 14,
+      });
     });
   });
 
@@ -104,8 +126,14 @@ describe('getPageBodyForContext', () => {
 
       // Expected: start=800, end=1300 (no shortfall needed)
       const expectedContent = longContent.slice(800, 1300);
-      expect(result).toBe(expectedContent);
-      expect(result).toHaveLength(500); // 1300 - 800 = 500
+      expect(result).toEqual({
+        content: expectedContent,
+        isPartial: true,
+        startIndex: 800,
+        endIndex: 1300,
+        totalLength: 2000,
+      });
+      expect(result?.content).toHaveLength(500); // 1300 - 800 = 500
     });
 
     it('should compensate shortfall when cursor is near document end', () => {
@@ -127,8 +155,14 @@ describe('getPageBodyForContext', () => {
       // Chars before: 100 + 150 = 250
       // Expected: start=max(0, 950-250)=700, end=950+50=1000
       const expectedContent = longContent.slice(700, 1000);
-      expect(result).toBe(expectedContent);
-      expect(result).toHaveLength(300); // 1000 - 700 = 300
+      expect(result).toEqual({
+        content: expectedContent,
+        isPartial: true,
+        startIndex: 700,
+        endIndex: 1000,
+        totalLength: 1000,
+      });
+      expect(result?.content).toHaveLength(300); // 1000 - 700 = 300
     });
 
     it('should handle extreme case: cursor at document end', () => {
@@ -150,8 +184,14 @@ describe('getPageBodyForContext', () => {
       // Chars before: 100 + 200 = 300
       // Expected: start=max(0, 1000-300)=700, end=1000+0=1000
       const expectedContent = longContent.slice(700, 1000);
-      expect(result).toBe(expectedContent);
-      expect(result).toHaveLength(300); // 1000 - 700 = 300
+      expect(result).toEqual({
+        content: expectedContent,
+        isPartial: true,
+        startIndex: 700,
+        endIndex: 1000,
+        totalLength: 1000,
+      });
+      expect(result?.content).toHaveLength(300); // 1000 - 700 = 300
     });
 
     it('should handle cursor at document start with startPos boundary', () => {
@@ -168,14 +208,19 @@ describe('getPageBodyForContext', () => {
 
       const result = getPageBodyForContext(mockEditor, 100, 200);
 
-      // Available after cursor: 1000
-      // Chars after: min(200, 1000) = 200
-      // Shortfall: 200 - 200 = 0
-      // Chars before: 100 + 0 = 100
-      // Expected: start=max(0, 0-100)=0, end=0+200=200
-      const expectedContent = longContent.slice(0, 200);
-      expect(result).toBe(expectedContent);
-      expect(result).toHaveLength(200);
+      // Available before cursor: 0 (cursor at start)
+      // Shortfall before: 100 - 0 = 100
+      // Chars after: 200 + 100 = 300
+      // Expected: start=0, end=0+300=300
+      const expectedContent = longContent.slice(0, 300);
+      expect(result).toEqual({
+        content: expectedContent,
+        isPartial: true,
+        startIndex: 0,
+        endIndex: 300,
+        totalLength: 1000,
+      });
+      expect(result?.content).toHaveLength(300);
     });
 
     it('should handle truly extreme shortfall with cursor very near end', () => {
@@ -197,8 +242,14 @@ describe('getPageBodyForContext', () => {
       // Chars before: 50 + 495 = 545
       // Expected: start=max(0, 995-545)=450, end=995+5=1000
       const expectedContent = longContent.slice(450, 1000);
-      expect(result).toBe(expectedContent);
-      expect(result).toHaveLength(550); // 1000 - 450 = 550
+      expect(result).toEqual({
+        content: expectedContent,
+        isPartial: true,
+        startIndex: 450,
+        endIndex: 1000,
+        totalLength: 1000,
+      });
+      expect(result?.content).toHaveLength(550); // 1000 - 450 = 550
     });
 
   });

+ 44 - 12
apps/app/src/features/openai/client/services/editor-assistant/get-page-body-for-context.ts

@@ -1,42 +1,74 @@
 import type { UseCodeMirrorEditor } from '@growi/editor/dist/client/services/use-codemirror-editor';
 
+export type PageBodyContextResult = {
+  content: string;
+  isPartial: boolean;
+  startIndex?: number; // Only present when partial
+  endIndex?: number; // Only present when partial
+  totalLength: number; // Total length of the original document
+};
+
 /**
  * Get page body text for AI context processing
  * @param codeMirrorEditor - CodeMirror editor instance
  * @param maxLengthBeforeCursor - Maximum number of characters to include before cursor position
  * @param maxLengthAfterCursor - Maximum number of characters to include after cursor position
- * @returns Page body text optimized for AI context, or undefined if editor is not available
+ * @returns Page body context result with metadata, or undefined if editor is not available
  */
 export const getPageBodyForContext = (
     codeMirrorEditor: UseCodeMirrorEditor | undefined,
     maxLengthBeforeCursor: number,
     maxLengthAfterCursor: number,
-): string | undefined => {
+): PageBodyContextResult | undefined => {
   const doc = codeMirrorEditor?.getDoc();
   const length = doc?.length ?? 0;
 
+  if (length === 0 || !doc) {
+    return undefined;
+  }
+
   const maxTotalLength = maxLengthBeforeCursor + maxLengthAfterCursor;
 
   if (length > maxTotalLength) {
     // Get cursor position
     const cursorPos = codeMirrorEditor?.view?.state.selection.main.head ?? 0;
 
-    // Calculate how many characters are available after cursor
+    // Calculate how many characters are available before and after cursor
+    const availableBeforeCursor = cursorPos;
     const availableAfterCursor = length - cursorPos;
+
+    // Calculate actual chars to take before and after cursor
+    const charsBeforeCursor = Math.min(maxLengthBeforeCursor, availableBeforeCursor);
     const charsAfterCursor = Math.min(maxLengthAfterCursor, availableAfterCursor);
 
-    // If chars after cursor is less than maxLengthAfterCursor, add the difference to chars before cursor
-    const shortfall = maxLengthAfterCursor - charsAfterCursor;
-    const charsBeforeCursor = maxLengthBeforeCursor + shortfall;
+    // Calculate shortfalls and redistribute
+    const shortfallBefore = maxLengthBeforeCursor - charsBeforeCursor;
+    const shortfallAfter = maxLengthAfterCursor - charsAfterCursor;
+
+    // Redistribute shortfalls
+    const finalCharsAfterCursor = Math.min(charsAfterCursor + shortfallBefore, availableAfterCursor);
+    const finalCharsBeforeCursor = Math.min(charsBeforeCursor + shortfallAfter, availableBeforeCursor);
 
-    // Calculate start position (cursor - charsBeforeCursor, but not less than 0)
-    const startPos = Math.max(0, cursorPos - charsBeforeCursor);
+    // Calculate start and end positions
+    const startPos = cursorPos - finalCharsBeforeCursor;
+    const endPos = cursorPos + finalCharsAfterCursor;
 
-    // Calculate end position
-    const endPos = cursorPos + charsAfterCursor;
+    const content = doc.slice(startPos, endPos).toString();
 
-    return doc?.slice(startPos, endPos).toString();
+    return {
+      content,
+      isPartial: true,
+      startIndex: startPos,
+      endIndex: endPos,
+      totalLength: length,
+    };
   }
 
-  return codeMirrorEditor?.getDocString();
+  const content = codeMirrorEditor?.getDocString() ?? '';
+
+  return {
+    content,
+    isPartial: false,
+    totalLength: length,
+  };
 };

+ 25 - 12
apps/app/src/features/openai/client/services/editor-assistant/use-editor-assistant.tsx

@@ -13,6 +13,10 @@ import { useTranslation } from 'react-i18next';
 import { type Text as YText } from 'yjs';
 
 import { apiv3Post } from '~/client/util/apiv3-client';
+import { useIsEnableUnifiedMergeView } from '~/stores-universal/context';
+import { useCurrentPageId } from '~/stores/page';
+
+import type { AiAssistantHasId } from '../../../interfaces/ai-assistant';
 import {
   SseMessageSchema,
   SseDetectedDiffSchema,
@@ -20,15 +24,12 @@ import {
   type SseMessage,
   type SseDetectedDiff,
   type SseFinalized,
-} from '~/features/openai/interfaces/editor-assistant/sse-schemas';
-import { handleIfSuccessfullyParsed } from '~/features/openai/utils/handle-if-successfully-parsed';
-import { useIsEnableUnifiedMergeView } from '~/stores-universal/context';
-import { useCurrentPageId } from '~/stores/page';
-
-import type { AiAssistantHasId } from '../../../interfaces/ai-assistant';
+  type EditRequestBody,
+} from '../../../interfaces/editor-assistant/sse-schemas';
 import type { MessageLog } from '../../../interfaces/message';
 import type { IThreadRelationHasId } from '../../../interfaces/thread-relation';
 import { ThreadType } from '../../../interfaces/thread-relation';
+import { handleIfSuccessfullyParsed } from '../../../utils/handle-if-successfully-parsed';
 import { AiAssistantDropdown } from '../../components/AiAssistant/AiAssistantSidebar/AiAssistantDropdown';
 import { QuickMenuList } from '../../components/AiAssistant/AiAssistantSidebar/QuickMenuList';
 import { useAiAssistantSidebar } from '../../stores/ai-assistant';
@@ -160,15 +161,27 @@ export const useEditorAssistant: UseEditorAssistant = () => {
     // Disable UnifiedMergeView when a Form is submitted with UnifiedMergeView enabled
     mutateIsEnableUnifiedMergeView(false);
 
+    const pageBodyContext = getPageBodyForContext(codeMirrorEditor, 2000, 8000);
+
+    if (!pageBodyContext) {
+      throw new Error('Unable to get page body context');
+    }
+
+    const requestBody = {
+      threadId,
+      userMessage: formData.input,
+      selectedText,
+      pageBody: pageBodyContext.content,
+      ...(pageBodyContext.isPartial && {
+        isPageBodyPartial: pageBodyContext.isPartial,
+        partialPageBodyStartIndex: pageBodyContext.startIndex,
+      }),
+    } satisfies EditRequestBody;
+
     const response = await fetch('/_api/v3/openai/edit', {
       method: 'POST',
       headers: { 'Content-Type': 'application/json' },
-      body: JSON.stringify({
-        threadId,
-        userMessage: formData.input,
-        selectedText,
-        pageBody: getPageBodyForContext(codeMirrorEditor, 2000, 8000),
-      }),
+      body: JSON.stringify(requestBody),
     });
 
     return response;