@lobehub/lobehub 2.0.0-next.360 → 2.0.0-next.362

package/packages/context-engine/src/engine/messages/__tests__/MessagesEngine.test.ts

@@ -407,4 +407,368 @@ describe('MessagesEngine', () => {
       expect(userMessage?.content).toBe('Please respond to: user input');
     });
   });
+
+  describe('Page Editor context', () => {
+    it('should inject page content to the last user message when pageContentContext is provided', async () => {
+      const messages: UIChatMessage[] = [
+        {
+          content: 'First question',
+          createdAt: Date.now(),
+          id: 'msg-1',
+          role: 'user',
+          updatedAt: Date.now(),
+        } as UIChatMessage,
+        {
+          content: 'Answer',
+          createdAt: Date.now(),
+          id: 'msg-2',
+          role: 'assistant',
+          updatedAt: Date.now(),
+        } as UIChatMessage,
+        {
+          content: 'Second question about the page',
+          createdAt: Date.now(),
+          id: 'msg-3',
+          role: 'user',
+          updatedAt: Date.now(),
+        } as UIChatMessage,
+      ];
+
+      const params = createBasicParams({
+        messages,
+        pageContentContext: {
+          markdown: '# Document Title\n\nDocument content here.',
+          metadata: {
+            charCount: 40,
+            lineCount: 3,
+            title: 'Test Document',
+          },
+        },
+      });
+      const engine = new MessagesEngine(params);
+
+      const result = await engine.process();
+
+      expect(result.messages).toEqual([
+        { content: 'First question', role: 'user' },
+        { content: 'Answer', role: 'assistant' },
+        {
+          content: `Second question about the page
+
+<!-- SYSTEM CONTEXT (NOT PART OF USER QUERY) -->
+<context.instruction>following part contains context information injected by the system. Please follow these instructions:
+
+1. Always prioritize handling user-visible content.
+2. the context is only required when user's queries rely on it.
+</context.instruction>
+<current_page_context>
+<current_page title="Test Document">
+<markdown chars="40" lines="3">
+# Document Title
+
+Document content here.
+</markdown>
+</current_page>
+</current_page_context>
+<!-- END SYSTEM CONTEXT -->`,
+          role: 'user',
+        },
+      ]);
+
+      expect(result.metadata.pageEditorContextInjected).toBe(true);
+    });
+
+    it('should not inject page content when not enabled', async () => {
+      const messages: UIChatMessage[] = [
+        {
+          content: 'Question',
+          createdAt: Date.now(),
+          id: 'msg-1',
+          role: 'user',
+          updatedAt: Date.now(),
+        } as UIChatMessage,
+      ];
+
+      const params = createBasicParams({ messages });
+      const engine = new MessagesEngine(params);
+
+      const result = await engine.process();
+
+      expect(result.messages).toEqual([{ content: 'Question', role: 'user' }]);
+      expect(result.metadata.pageEditorContextInjected).toBeUndefined();
+    });
+  });
+
+  describe('Page Selections', () => {
+    it('should inject page selections to each user message that has them', async () => {
+      const messages: UIChatMessage[] = [
+        {
+          content: 'First question with selection',
+          createdAt: Date.now(),
+          id: 'msg-1',
+          metadata: {
+            pageSelections: [
+              {
+                content: 'Selected paragraph 1',
+                id: 'sel-1',
+                pageId: 'page-1',
+                xml: '<p>Selected paragraph 1</p>',
+              },
+            ],
+          },
+          role: 'user',
+          updatedAt: Date.now(),
+        } as UIChatMessage,
+        {
+          content: 'Answer to first',
+          createdAt: Date.now(),
+          id: 'msg-2',
+          role: 'assistant',
+          updatedAt: Date.now(),
+        } as UIChatMessage,
+        {
+          content: 'Second question with different selection',
+          createdAt: Date.now(),
+          id: 'msg-3',
+          metadata: {
+            pageSelections: [
+              {
+                content: 'Selected paragraph 2',
+                id: 'sel-2',
+                pageId: 'page-1',
+                xml: '<p>Selected paragraph 2</p>',
+              },
+            ],
+          },
+          role: 'user',
+          updatedAt: Date.now(),
+        } as UIChatMessage,
+      ];
+
+      const params = createBasicParams({
+        messages,
+        pageContentContext: {
+          markdown: '# Doc',
+          metadata: { title: 'Doc' },
+        },
+      });
+      const engine = new MessagesEngine(params);
+
+      const result = await engine.process();
+
+      expect(result.messages).toEqual([
+        {
+          content: `First question with selection
+
+<!-- SYSTEM CONTEXT (NOT PART OF USER QUERY) -->
+<context.instruction>following part contains context information injected by the system. Please follow these instructions:
+
+1. Always prioritize handling user-visible content.
+2. the context is only required when user's queries rely on it.
+</context.instruction>
+<user_page_selections>
+<user_selections count="1">
+<selection >
+<p>Selected paragraph 1</p>
+</selection>
+</user_selections>
+</user_page_selections>
+<!-- END SYSTEM CONTEXT -->`,
+          role: 'user',
+        },
+        { content: 'Answer to first', role: 'assistant' },
+        {
+          content: `Second question with different selection
+
+<!-- SYSTEM CONTEXT (NOT PART OF USER QUERY) -->
+<context.instruction>following part contains context information injected by the system. Please follow these instructions:
+
+1. Always prioritize handling user-visible content.
+2. the context is only required when user's queries rely on it.
+</context.instruction>
+<user_page_selections>
+<user_selections count="1">
+<selection >
+<p>Selected paragraph 2</p>
+</selection>
+</user_selections>
+</user_page_selections>
+<current_page_context>
+<current_page title="Doc">
+<markdown chars="5" lines="1">
+# Doc
+</markdown>
+</current_page>
+</current_page_context>
+<!-- END SYSTEM CONTEXT -->`,
+          role: 'user',
+        },
+      ]);
+    });
+
+    it('should skip user messages without pageSelections', async () => {
+      const messages: UIChatMessage[] = [
+        {
+          content: 'No selection here',
+          createdAt: Date.now(),
+          id: 'msg-1',
+          role: 'user',
+          updatedAt: Date.now(),
+        } as UIChatMessage,
+        {
+          content: 'Answer',
+          createdAt: Date.now(),
+          id: 'msg-2',
+          role: 'assistant',
+          updatedAt: Date.now(),
+        } as UIChatMessage,
+        {
+          content: 'With selection',
+          createdAt: Date.now(),
+          id: 'msg-3',
+          metadata: {
+            pageSelections: [
+              {
+                content: 'Selected text',
+                id: 'sel-1',
+                pageId: 'page-1',
+                xml: '<span>Selected text</span>',
+              },
+            ],
+          },
+          role: 'user',
+          updatedAt: Date.now(),
+        } as UIChatMessage,
+      ];
+
+      const params = createBasicParams({
+        messages,
+        pageContentContext: {
+          markdown: '# Doc',
+          metadata: { title: 'Doc' },
+        },
+      });
+      const engine = new MessagesEngine(params);
+
+      const result = await engine.process();
+
+      expect(result.messages).toEqual([
+        { content: 'No selection here', role: 'user' },
+        { content: 'Answer', role: 'assistant' },
+        {
+          content: `With selection
+
+<!-- SYSTEM CONTEXT (NOT PART OF USER QUERY) -->
+<context.instruction>following part contains context information injected by the system. Please follow these instructions:
+
+1. Always prioritize handling user-visible content.
+2. the context is only required when user's queries rely on it.
+</context.instruction>
+<user_page_selections>
+<user_selections count="1">
+<selection >
+<span>Selected text</span>
+</selection>
+</user_selections>
+</user_page_selections>
+<current_page_context>
+<current_page title="Doc">
+<markdown chars="5" lines="1">
+# Doc
+</markdown>
+</current_page>
+</current_page_context>
+<!-- END SYSTEM CONTEXT -->`,
+          role: 'user',
+        },
+      ]);
+    });
+
+    it('should have only one SYSTEM CONTEXT wrapper when both selections and page content are injected', async () => {
+      const messages: UIChatMessage[] = [
+        {
+          content: 'Question about selection',
+          createdAt: Date.now(),
+          id: 'msg-1',
+          metadata: {
+            pageSelections: [
+              {
+                content: 'Selected text',
+                id: 'sel-1',
+                pageId: 'page-1',
+                xml: '<p>Selected text</p>',
+              },
+            ],
+          },
+          role: 'user',
+          updatedAt: Date.now(),
+        } as UIChatMessage,
+      ];
+
+      const params = createBasicParams({
+        messages,
+        pageContentContext: {
+          markdown: '# Full Document',
+          metadata: { title: 'Full Doc' },
+        },
+      });
+      const engine = new MessagesEngine(params);
+
+      const result = await engine.process();
+
+      expect(result.messages).toEqual([
+        {
+          content: `Question about selection
+
+<!-- SYSTEM CONTEXT (NOT PART OF USER QUERY) -->
+<context.instruction>following part contains context information injected by the system. Please follow these instructions:
+
+1. Always prioritize handling user-visible content.
+2. the context is only required when user's queries rely on it.
+</context.instruction>
+<user_page_selections>
+<user_selections count="1">
+<selection >
+<p>Selected text</p>
+</selection>
+</user_selections>
+</user_page_selections>
+<current_page_context>
+<current_page title="Full Doc">
+<markdown chars="15" lines="1">
+# Full Document
+</markdown>
+</current_page>
+</current_page_context>
+<!-- END SYSTEM CONTEXT -->`,
+          role: 'user',
+        },
+      ]);
+    });
+
+    it('should not inject selections when page editor is not enabled', async () => {
+      const messages: UIChatMessage[] = [
+        {
+          content: 'Question',
+          createdAt: Date.now(),
+          id: 'msg-1',
+          metadata: {
+            pageSelections: [
+              { content: 'Selected', id: 'sel-1', pageId: 'page-1', xml: '<p>Selected</p>' },
+            ],
+          },
+          role: 'user',
+          updatedAt: Date.now(),
+        } as UIChatMessage,
+      ];
+
+      // No pageContentContext or initialContext.pageEditor means not enabled
+      const params = createBasicParams({ messages });
+      const engine = new MessagesEngine(params);
+
+      const result = await engine.process();
+
+      expect(result.messages).toEqual([{ content: 'Question', role: 'user' }]);
+    });
+  });
 });

package/packages/context-engine/src/providers/PageEditorContextInjector.ts

@@ -20,6 +20,9 @@ export interface PageEditorContextInjectorConfig {
  * Page Editor Context Injector
  * Responsible for injecting current page context at the end of the last user message
  * This ensures the model receives the most up-to-date page/document state
+ *
+ * Note: Page selections (user-selected text regions) are handled separately by
+ * PageSelectionsInjector, which injects selections into each user message that has them
  */
 export class PageEditorContextInjector extends BaseLastUserContentProvider {
   readonly name = 'PageEditorContextInjector';
@@ -37,20 +40,11 @@ export class PageEditorContextInjector extends BaseLastUserContentProvider {
 
     const clonedContext = this.cloneContext(context);
 
-    // Skip if Page Editor is not enabled or no page content context
-    if (!this.config.enabled || !this.config.pageContentContext) {
-      log('Page Editor not enabled or no pageContentContext, skipping injection');
-      return this.markAsExecuted(clonedContext);
-    }
-
-    // Format page content context
-    const formattedContent = formatPageContentContext(this.config.pageContentContext);
-
-    log('Formatted content length:', formattedContent.length);
+    // Check if we have page content to inject
+    const hasPageContent = this.config.enabled && this.config.pageContentContext;
 
-    // Skip if no content to inject
-    if (!formattedContent) {
-      log('No content to inject after formatting');
+    if (!hasPageContent) {
+      log('No pageContentContext, skipping injection');
       return this.markAsExecuted(clonedContext);
     }
 
@@ -64,6 +58,16 @@ export class PageEditorContextInjector extends BaseLastUserContentProvider {
       return this.markAsExecuted(clonedContext);
     }
 
+    // Format page content
+    const formattedContent = formatPageContentContext(this.config.pageContentContext!);
+
+    if (!formattedContent) {
+      log('No content to inject after formatting');
+      return this.markAsExecuted(clonedContext);
+    }
+
+    log('Page content formatted, length:', formattedContent.length);
+
     // Check if system context wrapper already exists
     // If yes, only insert context block; if no, use full wrapper
     const hasExistingWrapper = this.hasExistingSystemContext(clonedContext);

package/packages/context-engine/src/providers/PageSelectionsInjector.ts

@@ -0,0 +1,65 @@
+import { formatPageSelections } from '@lobechat/prompts';
+import type { PageSelection } from '@lobechat/types';
+import debug from 'debug';
+
+import { BaseEveryUserContentProvider } from '../base/BaseEveryUserContentProvider';
+import type { Message, ProcessorOptions } from '../types';
+
+const log = debug('context-engine:provider:PageSelectionsInjector');
+
+export interface PageSelectionsInjectorConfig {
+  /** Whether Page Selections injection is enabled */
+  enabled?: boolean;
+}
+
+/**
+ * Page Selections Injector
+ * Responsible for injecting page selections into each user message that has them
+ * Unlike PageEditorContextInjector which only injects to the last user message,
+ * this processor handles selections attached to any user message in the conversation
+ *
+ * This injector runs BEFORE PageEditorContextInjector so that:
+ * - Each user message with selections gets a SYSTEM CONTEXT wrapper
+ * - PageEditorContextInjector can then reuse the wrapper for the last user message
+ */
+export class PageSelectionsInjector extends BaseEveryUserContentProvider {
+  readonly name = 'PageSelectionsInjector';
+
+  constructor(
+    private config: PageSelectionsInjectorConfig = {},
+    options: ProcessorOptions = {},
+  ) {
+    super(options);
+  }
+
+  protected buildContentForMessage(
+    message: Message,
+    index: number,
+  ): { content: string; contextType: string } | null {
+    // Skip if not enabled
+    if (!this.config.enabled) {
+      return null;
+    }
+
+    // Check if message has pageSelections in metadata
+    const pageSelections = message.metadata?.pageSelections as PageSelection[] | undefined;
+
+    if (!pageSelections || pageSelections.length === 0) {
+      return null;
+    }
+
+    // Format the selections
+    const formattedSelections = formatPageSelections(pageSelections);
+
+    if (!formattedSelections) {
+      return null;
+    }
+
+    log(`Building content for message at index ${index} with ${pageSelections.length} selections`);
+
+    return {
+      content: formattedSelections,
+      contextType: 'user_page_selections',
+    };
+  }
+}