@lobehub/lobehub 2.0.0-next.3 → 2.0.0-next.31

package/src/server/services/message/__tests__/index.test.ts

@@ -0,0 +1,348 @@
+import { LobeChatDatabase } from '@lobechat/database';
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { MessageModel } from '@/database/models/message';
+import { FileService } from '@/server/services/file';
+
+import { MessageService } from '../index';
+
+vi.mock('@/database/models/message');
+vi.mock('@/server/services/file');
+
+describe('MessageService', () => {
+  let messageService: MessageService;
+  let mockDB: LobeChatDatabase;
+  let mockMessageModel: MessageModel;
+  let mockFileService: FileService;
+  const userId = 'test-user-id';
+
+  beforeEach(() => {
+    mockDB = {} as LobeChatDatabase;
+    mockMessageModel = {
+      create: vi.fn(),
+      deleteMessage: vi.fn(),
+      deleteMessages: vi.fn(),
+      query: vi.fn(),
+      update: vi.fn(),
+      updateMessagePlugin: vi.fn(),
+      updateMessageRAG: vi.fn(),
+      updatePluginState: vi.fn(),
+    } as any;
+
+    mockFileService = {
+      getFullFileUrl: vi.fn().mockImplementation((path) => Promise.resolve(`/files${path}`)),
+    } as any;
+
+    // Mock constructors
+    vi.mocked(MessageModel).mockImplementation(() => mockMessageModel);
+    vi.mocked(FileService).mockImplementation(() => mockFileService);
+
+    messageService = new MessageService(mockDB, userId);
+  });
+
+  describe('removeMessage', () => {
+    it('should delete message and return { success: true } when no sessionId/topicId provided', async () => {
+      const messageId = 'msg-1';
+
+      const result = await messageService.removeMessage(messageId);
+
+      expect(mockMessageModel.deleteMessage).toHaveBeenCalledWith(messageId);
+      expect(result).toEqual({ success: true });
+      expect(mockMessageModel.query).not.toHaveBeenCalled();
+    });
+
+    it('should delete message and return message list when sessionId provided', async () => {
+      const messageId = 'msg-1';
+      const mockMessages = [{ id: 'msg-2', content: 'test' }];
+      vi.mocked(mockMessageModel.query).mockResolvedValue(mockMessages as any);
+
+      const result = await messageService.removeMessage(messageId, { sessionId: 'session-1' });
+
+      expect(mockMessageModel.deleteMessage).toHaveBeenCalledWith(messageId);
+      expect(mockMessageModel.query).toHaveBeenCalledWith(
+        { groupId: undefined, sessionId: 'session-1', topicId: undefined },
+        expect.objectContaining({
+          groupAssistantMessages: false,
+        }),
+      );
+      expect(result).toEqual({ messages: mockMessages, success: true });
+    });
+
+    it('should delete message and return message list when topicId provided', async () => {
+      const messageId = 'msg-1';
+      const mockMessages = [{ id: 'msg-2', content: 'test' }];
+      vi.mocked(mockMessageModel.query).mockResolvedValue(mockMessages as any);
+
+      const result = await messageService.removeMessage(messageId, { topicId: 'topic-1' });
+
+      expect(mockMessageModel.deleteMessage).toHaveBeenCalledWith(messageId);
+      expect(mockMessageModel.query).toHaveBeenCalledWith(
+        { groupId: undefined, sessionId: undefined, topicId: 'topic-1' },
+        expect.objectContaining({
+          groupAssistantMessages: false,
+        }),
+      );
+      expect(result).toEqual({ messages: mockMessages, success: true });
+    });
+  });
+
+  describe('removeMessages', () => {
+    it('should delete messages and return { success: true } when no sessionId/topicId provided', async () => {
+      const messageIds = ['msg-1', 'msg-2'];
+
+      const result = await messageService.removeMessages(messageIds);
+
+      expect(mockMessageModel.deleteMessages).toHaveBeenCalledWith(messageIds);
+      expect(result).toEqual({ success: true });
+      expect(mockMessageModel.query).not.toHaveBeenCalled();
+    });
+
+    it('should delete messages and return message list when sessionId provided', async () => {
+      const messageIds = ['msg-1', 'msg-2'];
+      const mockMessages = [{ id: 'msg-3', content: 'test' }];
+      vi.mocked(mockMessageModel.query).mockResolvedValue(mockMessages as any);
+
+      const result = await messageService.removeMessages(messageIds, { sessionId: 'session-1' });
+
+      expect(mockMessageModel.deleteMessages).toHaveBeenCalledWith(messageIds);
+      expect(mockMessageModel.query).toHaveBeenCalled();
+      expect(result).toEqual({ messages: mockMessages, success: true });
+    });
+  });
+
+  describe('updateMessageRAG', () => {
+    it('should update RAG and return { success: true } when no sessionId/topicId provided', async () => {
+      const messageId = 'msg-1';
+      const ragValue = { fileChunks: [{ id: 'chunk-1', similarity: 0.95 }] };
+
+      const result = await messageService.updateMessageRAG(messageId, ragValue);
+
+      expect(mockMessageModel.updateMessageRAG).toHaveBeenCalledWith(messageId, ragValue);
+      expect(result).toEqual({ success: true });
+      expect(mockMessageModel.query).not.toHaveBeenCalled();
+    });
+
+    it('should update RAG and return message list when sessionId provided', async () => {
+      const messageId = 'msg-1';
+      const ragValue = { fileChunks: [{ id: 'chunk-1', similarity: 0.95 }] };
+      const mockMessages = [{ id: 'msg-1', content: 'test' }];
+      vi.mocked(mockMessageModel.query).mockResolvedValue(mockMessages as any);
+
+      const result = await messageService.updateMessageRAG(messageId, ragValue, {
+        sessionId: 'session-1',
+      });
+
+      expect(mockMessageModel.updateMessageRAG).toHaveBeenCalledWith(messageId, ragValue);
+      expect(mockMessageModel.query).toHaveBeenCalled();
+      expect(result).toEqual({ messages: mockMessages, success: true });
+    });
+  });
+
+  describe('updatePluginError', () => {
+    it('should update plugin error and return { success: true } when no sessionId/topicId provided', async () => {
+      const messageId = 'msg-1';
+      const error = { type: 'TestError', message: 'Test error message' };
+
+      const result = await messageService.updatePluginError(messageId, error);
+
+      expect(mockMessageModel.updateMessagePlugin).toHaveBeenCalledWith(messageId, { error });
+      expect(result).toEqual({ success: true });
+      expect(mockMessageModel.query).not.toHaveBeenCalled();
+    });
+
+    it('should update plugin error and return message list when sessionId provided', async () => {
+      const messageId = 'msg-1';
+      const error = { type: 'TestError', message: 'Test error message' };
+      const mockMessages = [{ id: 'msg-1', content: 'test' }];
+      vi.mocked(mockMessageModel.query).mockResolvedValue(mockMessages as any);
+
+      const result = await messageService.updatePluginError(messageId, error, {
+        sessionId: 'session-1',
+      });
+
+      expect(mockMessageModel.updateMessagePlugin).toHaveBeenCalledWith(messageId, { error });
+      expect(mockMessageModel.query).toHaveBeenCalled();
+      expect(result).toEqual({ messages: mockMessages, success: true });
+    });
+  });
+
+  describe('updatePluginState', () => {
+    it('should update plugin state and return { success: true } when no sessionId/topicId provided', async () => {
+      const messageId = 'msg-1';
+      const state = { key: 'value' };
+
+      const result = await messageService.updatePluginState(messageId, state, {});
+
+      expect(mockMessageModel.updatePluginState).toHaveBeenCalledWith(messageId, state);
+      expect(result).toEqual({ success: true });
+      expect(mockMessageModel.query).not.toHaveBeenCalled();
+    });
+
+    it('should update plugin state and return message list when sessionId provided', async () => {
+      const messageId = 'msg-1';
+      const state = { key: 'value' };
+      const mockMessages = [{ id: 'msg-1', content: 'test' }];
+      vi.mocked(mockMessageModel.query).mockResolvedValue(mockMessages as any);
+
+      const result = await messageService.updatePluginState(messageId, state, {
+        sessionId: 'session-1',
+      });
+
+      expect(mockMessageModel.updatePluginState).toHaveBeenCalledWith(messageId, state);
+      expect(mockMessageModel.query).toHaveBeenCalled();
+      expect(result).toEqual({ messages: mockMessages, success: true });
+    });
+  });
+
+  describe('updateMessage', () => {
+    it('should update message and return { success: true } when no sessionId/topicId provided', async () => {
+      const messageId = 'msg-1';
+      const value = { content: 'updated content' };
+
+      const result = await messageService.updateMessage(messageId, value as any, {});
+
+      expect(mockMessageModel.update).toHaveBeenCalledWith(messageId, value);
+      expect(result).toEqual({ success: true });
+      expect(mockMessageModel.query).not.toHaveBeenCalled();
+    });
+
+    it('should update message and return message list when sessionId provided', async () => {
+      const messageId = 'msg-1';
+      const value = { content: 'updated content' };
+      const mockMessages = [{ id: 'msg-1', content: 'updated content' }];
+      vi.mocked(mockMessageModel.query).mockResolvedValue(mockMessages as any);
+
+      const result = await messageService.updateMessage(messageId, value as any, {
+        sessionId: 'session-1',
+      });
+
+      expect(mockMessageModel.update).toHaveBeenCalledWith(messageId, value);
+      expect(mockMessageModel.query).toHaveBeenCalled();
+      expect(result).toEqual({ messages: mockMessages, success: true });
+    });
+  });
+
+  describe('useGroup option', () => {
+    it('should pass useGroup option to query', async () => {
+      const mockMessages = [{ id: 'msg-1', content: 'test' }];
+      vi.mocked(mockMessageModel.query).mockResolvedValue(mockMessages as any);
+
+      await messageService.removeMessage('msg-1', {
+        sessionId: 'session-1',
+        useGroup: true,
+      });
+
+      expect(mockMessageModel.query).toHaveBeenCalledWith(
+        expect.anything(),
+        expect.objectContaining({
+          groupAssistantMessages: true,
+        }),
+      );
+    });
+
+    it('should default useGroup to false', async () => {
+      const mockMessages = [{ id: 'msg-1', content: 'test' }];
+      vi.mocked(mockMessageModel.query).mockResolvedValue(mockMessages as any);
+
+      await messageService.removeMessage('msg-1', { sessionId: 'session-1' });
+
+      expect(mockMessageModel.query).toHaveBeenCalledWith(
+        expect.anything(),
+        expect.objectContaining({
+          groupAssistantMessages: false,
+        }),
+      );
+    });
+  });
+
+  describe('createNewMessage', () => {
+    it('should create message and return message list', async () => {
+      const params = {
+        content: 'Hello',
+        role: 'user' as const,
+        sessionId: 'session-1',
+      };
+      const createdMessage = { id: 'msg-1', ...params };
+      const mockMessages = [createdMessage, { id: 'msg-2', content: 'Hi' }];
+
+      vi.mocked(mockMessageModel.create).mockResolvedValue(createdMessage as any);
+      vi.mocked(mockMessageModel.query).mockResolvedValue(mockMessages as any);
+
+      const result = await messageService.createNewMessage(params as any);
+
+      expect(mockMessageModel.create).toHaveBeenCalledWith(params);
+      expect(mockMessageModel.query).toHaveBeenCalledWith(
+        {
+          current: 0,
+          groupId: undefined,
+          pageSize: 9999,
+          sessionId: 'session-1',
+          topicId: undefined,
+        },
+        expect.objectContaining({
+          groupAssistantMessages: false,
+        }),
+      );
+      expect(result).toEqual({
+        id: 'msg-1',
+        messages: mockMessages,
+      });
+    });
+
+    it('should create message with useGroup option', async () => {
+      const params = {
+        content: 'Hello',
+        role: 'assistant' as const,
+        sessionId: 'session-1',
+      };
+      const createdMessage = { id: 'msg-1', ...params };
+      const mockMessages = [createdMessage];
+
+      vi.mocked(mockMessageModel.create).mockResolvedValue(createdMessage as any);
+      vi.mocked(mockMessageModel.query).mockResolvedValue(mockMessages as any);
+
+      const result = await messageService.createNewMessage(params as any, { useGroup: true });
+
+      expect(mockMessageModel.query).toHaveBeenCalledWith(
+        expect.anything(),
+        expect.objectContaining({
+          groupAssistantMessages: true,
+        }),
+      );
+      expect(result).toEqual({
+        id: 'msg-1',
+        messages: mockMessages,
+      });
+    });
+
+    it('should create message with topicId and groupId', async () => {
+      const params = {
+        content: 'Hello',
+        groupId: 'group-1',
+        role: 'user' as const,
+        sessionId: 'session-1',
+        topicId: 'topic-1',
+      };
+      const createdMessage = { id: 'msg-1', ...params };
+      const mockMessages = [createdMessage];
+
+      vi.mocked(mockMessageModel.create).mockResolvedValue(createdMessage as any);
+      vi.mocked(mockMessageModel.query).mockResolvedValue(mockMessages as any);
+
+      const result = await messageService.createNewMessage(params as any);
+
+      expect(mockMessageModel.query).toHaveBeenCalledWith(
+        {
+          current: 0,
+          groupId: 'group-1',
+          pageSize: 9999,
+          sessionId: 'session-1',
+          topicId: 'topic-1',
+        },
+        expect.anything(),
+      );
+      expect(result.id).toBe('msg-1');
+      expect(result.messages).toEqual(mockMessages);
+    });
+  });
+});

package/src/server/services/message/index.ts

@@ -0,0 +1,159 @@
+import { LobeChatDatabase } from '@lobechat/database';
+import { CreateMessageParams, UpdateMessageParams } from '@lobechat/types';
+
+import { MessageModel } from '@/database/models/message';
+
+import { FileService } from '../file';
+
+interface QueryOptions {
+  groupId?: string | null;
+  sessionId?: string | null;
+  topicId?: string | null;
+  useGroup?: boolean;
+}
+
+interface CreateMessageResult {
+  id: string;
+  messages: any[];
+}
+
+/**
+ * Message Service
+ *
+ * Encapsulates repeated "mutation + conditional query" logic.
+ * After performing update/delete operations, conditionally returns message list based on sessionId/topicId.
+ */
+export class MessageService {
+  private messageModel: MessageModel;
+  private fileService: FileService;
+
+  constructor(db: LobeChatDatabase, userId: string) {
+    this.messageModel = new MessageModel(db, userId);
+    this.fileService = new FileService(db, userId);
+  }
+
+  /**
+   * Unified URL processing function
+   */
+  private get postProcessUrl() {
+    return (path: string | null) => this.fileService.getFullFileUrl(path);
+  }
+
+  /**
+   * Unified query options
+   */
+  private getQueryOptions(options: QueryOptions) {
+    return {
+      groupAssistantMessages: options.useGroup ?? false,
+      postProcessUrl: this.postProcessUrl,
+    };
+  }
+
+  /**
+   * Query messages and return response with success status (used after mutations)
+   */
+  private async queryWithSuccess(options?: QueryOptions) {
+    if (!options || (options.sessionId === undefined && options.topicId === undefined)) {
+      return { success: true };
+    }
+
+    const { sessionId, topicId, groupId } = options;
+
+    const messages = await this.messageModel.query(
+      { groupId, sessionId, topicId },
+      this.getQueryOptions(options),
+    );
+
+    return { messages, success: true };
+  }
+
+  /**
+   * Remove messages with optional message list return
+   * Pattern: delete + conditional query
+   */
+  async removeMessages(ids: string[], options?: QueryOptions) {
+    await this.messageModel.deleteMessages(ids);
+    return this.queryWithSuccess(options);
+  }
+
+  /**
+   * Remove single message with optional message list return
+   * Pattern: delete + conditional query
+   */
+  async removeMessage(id: string, options?: QueryOptions) {
+    await this.messageModel.deleteMessage(id);
+    return this.queryWithSuccess(options);
+  }
+
+  /**
+   * Update message RAG with optional message list return
+   * Pattern: update + conditional query
+   */
+  async updateMessageRAG(id: string, value: any, options?: QueryOptions) {
+    await this.messageModel.updateMessageRAG(id, value);
+    return this.queryWithSuccess(options);
+  }
+
+  /**
+   * Update plugin error with optional message list return
+   * Pattern: update + conditional query
+   */
+  async updatePluginError(id: string, value: any, options?: QueryOptions) {
+    await this.messageModel.updateMessagePlugin(id, { error: value });
+    return this.queryWithSuccess(options);
+  }
+
+  /**
+   * Update plugin state and return message list
+   * Pattern: update + conditional query
+   */
+  async updatePluginState(id: string, value: any, options: QueryOptions): Promise<any> {
+    await this.messageModel.updatePluginState(id, value);
+    return this.queryWithSuccess(options);
+  }
+
+  /**
+   * Update message and return message list
+   * Pattern: update + conditional query
+   */
+  async updateMessage(id: string, value: UpdateMessageParams, options: QueryOptions): Promise<any> {
+    await this.messageModel.update(id, value as any);
+    return this.queryWithSuccess(options);
+  }
+
+  /**
+   * Create a new message and return the complete message list
+   * Pattern: create + query
+   *
+   * This method combines message creation and querying into a single operation,
+   * reducing the need for separate refresh calls and improving performance.
+   */
+  async createNewMessage(
+    params: CreateMessageParams,
+    options?: QueryOptions,
+  ): Promise<CreateMessageResult> {
+    // 1. Create the message
+    const item = await this.messageModel.create(params);
+
+    // 2. Query all messages for this session/topic
+    const messages = await this.messageModel.query(
+      {
+        current: 0,
+        groupId: params.groupId,
+        pageSize: 9999,
+        sessionId: params.sessionId,
+        topicId: params.topicId,
+      },
+      {
+        groupAssistantMessages: options?.useGroup ?? false,
+        postProcessUrl: this.postProcessUrl,
+      },
+    );
+
+    // 3. Return the result
+    return {
+      id: item.id,
+      messages,
+    };
+  }
+}

package/src/server/services/search/impls/firecrawl/index.ts

@@ -26,7 +26,7 @@ export class FirecrawlImpl implements SearchServiceImpl {
 
   private get baseUrl(): string {
     // Assuming the base URL is consistent with the crawl endpoint
-    return process.env.FIRECRAWL_URL || 'https://api.firecrawl.dev/v1';
+    return process.env.FIRECRAWL_URL || 'https://api.firecrawl.dev/v2';
   }
 
   async query(query: string, params: SearchParams = {}): Promise<UniformSearchResponse> {
@@ -34,13 +34,14 @@ export class FirecrawlImpl implements SearchServiceImpl {
     const endpoint = urlJoin(this.baseUrl, '/search');
 
     const defaultQueryParams: FirecrawlSearchParameters = {
-      limit: 15,
+      limit: 20,
       query,
       /*
       scrapeOptions: {
         formats: ["markdown"]
       },
       */
+      sources: [{ type: 'web' }, { type: 'news' }],
     };
 
     let body: FirecrawlSearchParameters = {
@@ -95,25 +96,64 @@ export class FirecrawlImpl implements SearchServiceImpl {
 
       log('Parsed Firecrawl response: %o', firecrawlResponse);
 
-      const mappedResults = (firecrawlResponse.data || []).map(
+      // V2 API returns data as object with web/images/news arrays
+      const webResults = firecrawlResponse.data.web || [];
+      const imageResults = firecrawlResponse.data.images || [];
+      const newsResults = firecrawlResponse.data.news || [];
+
+      // Map web results
+      const mappedWebResults = webResults.map(
+        (result): UniformSearchResult => ({
+          category: 'general',
+          content: result.description || result.markdown || '',
+          engines: ['firecrawl'],
+          parsedUrl: result.url ? new URL(result.url).hostname : '',
+          score: 1,
+          title: result.title || '',
+          url: result.url,
+        }),
+      );
+
+      // Map news results
+      const mappedNewsResults = newsResults.map(
         (result): UniformSearchResult => ({
-          category: 'general', // Default category
-          content: result.description || '', // Prioritize content, fallback to snippet
-          engines: ['firecrawl'], // Use 'firecrawl' as the engine name
-          parsedUrl: result.url ? new URL(result.url).hostname : '', // Basic URL parsing
-          score: 1, // Default score to 1
+          category: 'news',
+          content: result.snippet || result.markdown || '',
+          engines: ['firecrawl'],
+          parsedUrl: result.url ? new URL(result.url).hostname : '',
+          score: 1,
           title: result.title || '',
           url: result.url,
         }),
       );
 
-      log('Mapped %d results to SearchResult format', mappedResults.length);
+      // Map image results
+      const mappedImageResults = imageResults.map(
+        (result): UniformSearchResult => ({
+          category: 'images',
+          content: result.title || '',
+          engines: ['firecrawl'],
+          parsedUrl: result.url ? new URL(result.url).hostname : '',
+          score: 1,
+          title: result.title || '',
+          url: result.imageUrl, // Use imageUrl for images
+        }),
+      );
+
+      // Combine all results
+      const allResults = [...mappedWebResults, ...mappedNewsResults, ...mappedImageResults];
+
+      log('Mapped %d results to SearchResult format', allResults.length);
+
+      if (firecrawlResponse.warning) {
+        log.extend('warn')('Firecrawl warning: %s', firecrawlResponse.warning);
+      }
 
       return {
         costTime,
         query: query,
-        resultNumbers: mappedResults.length,
-        results: mappedResults,
+        resultNumbers: allResults.length,
+        results: allResults,
       };
     } catch (error) {
       log.extend('error')('Error parsing Firecrawl response: %o', error);

package/src/server/services/search/impls/firecrawl/type.ts

@@ -1,35 +1,86 @@
+// V2 API Types
 interface FirecrawlScrapeOptions {
-  formats: string[];
+  blockAds?: boolean;
+  formats?: string[];
+  maxAge?: number;
+  onlyMainContent?: boolean;
+  removeBase64Images?: boolean;
 }
 
+type FirecrawlSource =
+  | { location?: string; tbs?: string; type: 'web' }
+  | { type: 'images' }
+  | { type: 'news' };
+
+type FirecrawlCategory = { type: 'github' } | { type: 'research' } | { type: 'pdf' };
+
 export interface FirecrawlSearchParameters {
+  categories?: FirecrawlCategory[];
   country?: string;
-  lang?: string;
+  ignoreInvalidURLs?: boolean;
   limit?: number;
+  location?: string;
   query: string;
   scrapeOptions?: FirecrawlScrapeOptions;
+  sources?: FirecrawlSource[];
   tbs?: string;
   timeout?: number;
 }
 
 interface FirecrawlMetadata {
   description?: string;
+  error?: string | null;
   sourceURL?: string;
   statusCode?: number;
+  title?: string;
+}
+
+// Web search result
+interface FirecrawlWebResult {
+  description: string;
+  html?: string | null;
+  links?: string[];
+  markdown?: string | null;
+  metadata?: FirecrawlMetadata;
+  rawHtml?: string | null;
+  screenshot?: string | null;
   title: string;
+  url: string;
 }
 
-interface FirecrawlData {
-  description?: string;
-  html?: string;
+// Image search result
+interface FirecrawlImageResult {
+  imageHeight: number;
+  imageUrl: string;
+  imageWidth: number;
+  position: number;
+  title: string;
+  url: string;
+}
+
+// News search result
+interface FirecrawlNewsResult {
+  date: string;
+  html?: string | null;
+  imageUrl?: string;
   links?: string[];
-  markdown?: string;
+  markdown?: string | null;
   metadata?: FirecrawlMetadata;
-  title?: string;
+  position: number;
+  rawHtml?: string | null;
+  screenshot?: string | null;
+  snippet: string;
+  title: string;
   url: string;
 }
 
+// V2 Response structure
 export interface FirecrawlResponse {
-  data: FirecrawlData[];
-  success?: boolean;
+  data: {
+    images?: FirecrawlImageResult[];
+    news?: FirecrawlNewsResult[];
+    web?: FirecrawlWebResult[];
+  };
+  success: boolean;
+  warning?: string | null;
 }