@lobehub/chat 1.26.9 → 1.26.11

package/CHANGELOG.md

@@ -2,6 +2,48 @@
 
 # Changelog
 
+### [Version 1.26.11](https://github.com/lobehub/lobe-chat/compare/v1.26.10...v1.26.11)
+
+<sup>Released on **2024-10-29**</sup>
+
+<br/>
+
+<details>
+<summary><kbd>Improvements and Fixes</kbd></summary>
+
+</details>
+
+<div align="right">
+
+[![](https://img.shields.io/badge/-BACK_TO_TOP-151515?style=flat-square)](#readme-top)
+
+</div>
+
+### [Version 1.26.10](https://github.com/lobehub/lobe-chat/compare/v1.26.9...v1.26.10)
+
+<sup>Released on **2024-10-29**</sup>
+
+#### ♻ Code Refactoring
+
+- **misc**: Refactor the aiChat slice actions.
+
+<br/>
+
+<details>
+<summary><kbd>Improvements and Fixes</kbd></summary>
+
+#### Code refactoring
+
+- **misc**: Refactor the aiChat slice actions, closes [#4542](https://github.com/lobehub/lobe-chat/issues/4542) ([8b33ba4](https://github.com/lobehub/lobe-chat/commit/8b33ba4))
+
+</details>
+
+<div align="right">
+
+[![](https://img.shields.io/badge/-BACK_TO_TOP-151515?style=flat-square)](#readme-top)
+
+</div>
+
 ### [Version 1.26.9](https://github.com/lobehub/lobe-chat/compare/v1.26.8...v1.26.9)
 
 <sup>Released on **2024-10-29**</sup>

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@lobehub/chat",
-  "version": "1.26.9",
+  "version": "1.26.11",
   "description": "Lobe Chat - an open-source, high-performance chatbot framework that supports speech synthesis, multimodal, and extensible Function Call plugin system. Supports one-click free deployment of your private ChatGPT/LLM web application.",
   "keywords": [
     "framework",

package/src/database/server/models/chunk.ts

@@ -3,7 +3,7 @@ import { and, desc, isNull } from 'drizzle-orm/expressions';
 import { chunk } from 'lodash-es';
 
 import { serverDB } from '@/database/server';
-import { ChunkMetadata, FileChunk, SemanticSearchChunk } from '@/types/chunk';
+import { ChunkMetadata, FileChunk } from '@/types/chunk';
 
 import {
   NewChunkItem,
@@ -148,6 +148,8 @@ export class ChunkModel {
 
     const data = await serverDB
       .select({
+        fileId: fileChunks.fileId,
+        fileName: files.name,
         id: chunks.id,
         index: chunks.index,
         metadata: chunks.metadata,
@@ -158,16 +160,15 @@ export class ChunkModel {
       .from(chunks)
       .leftJoin(embeddings, eq(chunks.id, embeddings.chunkId))
       .leftJoin(fileChunks, eq(chunks.id, fileChunks.chunkId))
+      .leftJoin(files, eq(fileChunks.fileId, files.id))
       .where(fileIds ? inArray(fileChunks.fileId, fileIds) : undefined)
       .orderBy((t) => desc(t.similarity))
       .limit(30);
 
-    return data.map(
-      (item): SemanticSearchChunk => ({
-        ...item,
-        metadata: item.metadata as ChunkMetadata,
-      }),
-    );
+    return data.map((item) => ({
+      ...item,
+      metadata: item.metadata as ChunkMetadata,
+    }));
   }
 
   async semanticSearchForChat({
@@ -187,7 +188,7 @@ export class ChunkModel {
     const result = await serverDB
       .select({
         fileId: files.id,
-        filename: files.name,
+        fileName: files.name,
         id: chunks.id,
         index: chunks.index,
         metadata: chunks.metadata,
@@ -205,6 +206,8 @@ export class ChunkModel {
 
     return result.map((item) => {
       return {
+        fileId: item.fileId,
+        fileName: item.fileName,
         id: item.id,
         index: item.index,
         similarity: item.similarity,

package/src/features/ChatInput/useSend.ts

@@ -2,8 +2,8 @@ import { useCallback, useMemo } from 'react';
 
 import { useChatStore } from '@/store/chat';
 import { chatSelectors } from '@/store/chat/selectors';
-import { SendMessageParams } from '@/store/chat/slices/aiChat/action';
 import { fileChatSelectors, useFileStore } from '@/store/file';
+import { SendMessageParams } from '@/types/message';
 
 export type UseSendMessageParams = Pick<
   SendMessageParams,

package/src/prompts/files/index.test.ts

@@ -0,0 +1,117 @@
+import { describe, expect, it } from 'vitest';
+
+import { ChatFileItem, ChatImageItem } from '@/types/message';
+
+import { filesPrompts } from './index';
+
+describe('filesPrompts', () => {
+  // 创建测试用的示例数据
+  const mockImage: ChatImageItem = {
+    id: 'img-1',
+    alt: 'test image',
+    url: 'https://example.com/image.jpg',
+  };
+
+  const mockFile: ChatFileItem = {
+    id: 'file-1',
+    name: 'test.pdf',
+    fileType: 'application/pdf',
+    size: 1024,
+    url: 'https://example.com/test.pdf',
+  };
+
+  it('should generate prompt with only images', () => {
+    const result = filesPrompts({
+      imageList: [mockImage],
+      fileList: undefined,
+    });
+
+    expect(result).toEqual(
+      `<files_info>
+<images>
+<images_docstring>here are user upload images you can refer to</images_docstring>
+<image name="test image" url="https://example.com/image.jpg"></image>
+</images>
+
+</files_info>`,
+    );
+  });
+
+  it('should generate prompt with only files', () => {
+    const result = filesPrompts({
+      imageList: [],
+      fileList: [mockFile],
+    });
+
+    expect(result).toEqual(
+      `<files_info>
+
+<files>
+<files_docstring>here are user upload files you can refer to</files_docstring>
+<file id="file-1" name="test.pdf" type="application/pdf" size="1024" url="https://example.com/test.pdf"></file>
+</files>
+</files_info>`,
+    );
+  });
+
+  it('should generate prompt with both images and files', () => {
+    const result = filesPrompts({
+      imageList: [mockImage],
+      fileList: [mockFile],
+    });
+
+    expect(result).toEqual(
+      `<files_info>
+<images>
+<images_docstring>here are user upload images you can refer to</images_docstring>
+<image name="test image" url="https://example.com/image.jpg"></image>
+</images>
+<files>
+<files_docstring>here are user upload files you can refer to</files_docstring>
+<file id="file-1" name="test.pdf" type="application/pdf" size="1024" url="https://example.com/test.pdf"></file>
+</files>
+</files_info>`,
+    );
+  });
+
+  it('should generate prompt with empty lists', () => {
+    const result = filesPrompts({
+      imageList: [],
+      fileList: [],
+    });
+
+    expect(result).toEqual('');
+  });
+
+  it('should handle multiple images and files', () => {
+    const images: ChatImageItem[] = [
+      mockImage,
+      {
+        id: 'img-2',
+        alt: 'second image',
+        url: 'https://example.com/image2.jpg',
+      },
+    ];
+
+    const files: ChatFileItem[] = [
+      mockFile,
+      {
+        id: 'file-2',
+        name: 'document.docx',
+        fileType: 'application/docx',
+        size: 2048,
+        url: 'https://example.com/document.docx',
+      },
+    ];
+
+    const result = filesPrompts({
+      imageList: images,
+      fileList: files,
+    });
+
+    expect(result).toContain('second image');
+    expect(result).toContain('document.docx');
+    expect(result).toMatch(/<image.*?>.*<image.*?>/s); // Check for multiple image tags
+    expect(result).toMatch(/<file.*?>.*<file.*?>/s); // Check for multiple file tags
+  });
+});

package/src/prompts/files/index.ts

@@ -10,6 +10,8 @@ export const filesPrompts = ({
   fileList?: ChatFileItem[];
   imageList: ChatImageItem[];
 }) => {
+  if (imageList.length === 0 && (fileList || []).length === 0) return '';
+
   const prompt = `<files_info>
 ${imagesPrompts(imageList)}
 ${fileList ? filePrompts(fileList) : ''}

package/src/prompts/knowledgeBaseQA/__snapshots__/index.test.ts.snap

@@ -0,0 +1,26 @@
+// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
+
+exports[`knowledgeBaseQAPrompts > should generate prompt with all parameters 1`] = `
+"<knowledge_base_qa_info>
+You are also a helpful assistant good answering questions related to Test Knowledge. And you'll be provided with a question and several passages that might be relevant. And currently your task is to provide answer based on the question and passages.
+<knowledge_base_anwser_instruction>
+- Note that passages might not be relevant to the question, please only use the passages that are relevant.
+- if there is no relevant passage, please answer using your knowledge.
+- Answer should use the same original language as the question and follow markdown syntax.
+</knowledge_base_anwser_instruction>
+<knowledge_bases>
+<knowledge_bases_docstring>here are the knowledge base scope we retrieve chunks from:</knowledge_bases_docstring>
+<knowledge id="kb1" name="Test Knowledge" type="file" fileType="txt" >Test description</knowledge>
+</knowledge_bases>
+<retrieved_chunks>
+<retrieved_chunks_docstring>here are retrived chunks you can refer to:</retrieved_chunks_docstring>
+<chunk fileId="file1" fileName="test.txt" similarity="0.8"  pageNumber="1" >This is a test chunk</chunk>
+</retrieved_chunks>
+<user_query>
+<user_query_docstring>to make result better, we may rewrite user's question.If there is a rewrite query, it will be wrapper with \`rewrite_query\` tag.</user_query_docstring>
+
+<raw_query>What is the test about?</raw_query>
+<rewrite_query>Could you explain the content of the test?</rewrite_query>
+<user_query>
+</knowledge_base_qa_info>"
+`;

package/src/prompts/knowledgeBaseQA/chunk.ts

@@ -0,0 +1,15 @@
+import { ChatSemanticSearchChunk } from '@/types/chunk';
+
+const chunkPrompt = (item: ChatSemanticSearchChunk) =>
+  `<chunk fileId="${item.fileId}" fileName="${item.fileName}" similarity="${item.similarity}" ${item.pageNumber ? ` pageNumber="${item.pageNumber}" ` : ''}>${item.text}</chunk>`;
+
+export const chunkPrompts = (fileList: ChatSemanticSearchChunk[]) => {
+  if (fileList.length === 0) return '';
+
+  const prompt = `<retrieved_chunks>
+<retrieved_chunks_docstring>here are retrived chunks you can refer to:</retrieved_chunks_docstring>
+${fileList.map((item) => chunkPrompt(item)).join('\n')}
+</retrieved_chunks>`;
+
+  return prompt.trim();
+};

package/src/prompts/knowledgeBaseQA/index.test.ts

@@ -0,0 +1,146 @@
+import { describe, expect, it } from 'vitest';
+
+import { ChatSemanticSearchChunk } from '@/types/chunk';
+import { KnowledgeItem, KnowledgeType } from '@/types/knowledgeBase';
+
+import { knowledgeBaseQAPrompts } from './index';
+
+describe('knowledgeBaseQAPrompts', () => {
+  // Define test data
+  const mockChunks: ChatSemanticSearchChunk[] = [
+    {
+      id: '1',
+      fileId: 'file1',
+      fileName: 'test.txt',
+      text: 'This is a test chunk',
+      similarity: 0.8,
+      pageNumber: 1,
+    },
+  ];
+
+  const mockKnowledge: KnowledgeItem[] = [
+    {
+      id: 'kb1',
+      name: 'Test Knowledge',
+      type: KnowledgeType.File,
+      fileType: 'txt',
+      description: 'Test description',
+    },
+  ];
+
+  const userQuery = 'What is the test about?';
+  const rewriteQuery = 'Could you explain the content of the test?';
+
+  it('should return empty string if chunks is empty', () => {
+    const result = knowledgeBaseQAPrompts({
+      chunks: [],
+      knowledge: mockKnowledge,
+      userQuery,
+    });
+
+    expect(result).toBe('');
+  });
+
+  it('should return empty string if chunks is undefined', () => {
+    const result = knowledgeBaseQAPrompts({
+      knowledge: mockKnowledge,
+      userQuery,
+    });
+
+    expect(result).toBe('');
+  });
+
+  it('should generate prompt with all parameters', () => {
+    const result = knowledgeBaseQAPrompts({
+      chunks: mockChunks,
+      knowledge: mockKnowledge,
+      userQuery,
+      rewriteQuery,
+    });
+
+    // Verify the prompt structure and content
+    expect(result).toMatchSnapshot();
+  });
+
+  it('should generate prompt without rewriteQuery', () => {
+    const result = knowledgeBaseQAPrompts({
+      chunks: mockChunks,
+      knowledge: mockKnowledge,
+      userQuery,
+    });
+
+    expect(result).toContain('<raw_query>What is the test about?</raw_query>');
+    expect(result).not.toContain('<rewrite_query>');
+  });
+
+  it('should generate prompt without knowledge', () => {
+    const result = knowledgeBaseQAPrompts({
+      chunks: mockChunks,
+      userQuery,
+    });
+
+    expect(result).toContain(
+      'You are also a helpful assistant good answering questions related to',
+    );
+    expect(result).not.toContain('<knowledge_bases>');
+  });
+
+  it('should handle empty knowledge array', () => {
+    const result = knowledgeBaseQAPrompts({
+      chunks: mockChunks,
+      knowledge: [],
+      userQuery,
+    });
+
+    expect(result).toContain(
+      'You are also a helpful assistant good answering questions related to',
+    );
+    expect(result).not.toContain('<knowledge_bases>');
+  });
+
+  it('should properly escape special characters in input', () => {
+    const specialChunks: ChatSemanticSearchChunk[] = [
+      {
+        id: '1',
+        fileId: 'file1',
+        fileName: 'test&.txt',
+        text: 'This is a test with & < > "quotes"',
+        similarity: 0.8,
+      },
+    ];
+
+    const result = knowledgeBaseQAPrompts({
+      chunks: specialChunks,
+      userQuery: 'Test with & < > "quotes"',
+    });
+
+    expect(result).toContain('test&.txt');
+    expect(result).toContain('This is a test with & < > "quotes"');
+    expect(result).toContain('Test with & < > "quotes"');
+  });
+
+  it('should handle multiple knowledge items', () => {
+    const multipleKnowledge: KnowledgeItem[] = [
+      {
+        id: 'kb1',
+        name: 'Knowledge 1',
+        type: KnowledgeType.File,
+      },
+      {
+        id: 'kb2',
+        name: 'Knowledge 2',
+        type: KnowledgeType.KnowledgeBase,
+      },
+    ];
+
+    const result = knowledgeBaseQAPrompts({
+      chunks: mockChunks,
+      knowledge: multipleKnowledge,
+      userQuery,
+    });
+
+    expect(result).toContain('Knowledge 1/Knowledge 2');
+    expect(result).toContain('<knowledge id="kb1"');
+    expect(result).toContain('<knowledge id="kb2"');
+  });
+});

package/src/prompts/knowledgeBaseQA/index.ts

@@ -0,0 +1,33 @@
+import { chunkPrompts } from '@/prompts/knowledgeBaseQA/chunk';
+import { knowledgePrompts } from '@/prompts/knowledgeBaseQA/knowledge';
+import { userQueryPrompt } from '@/prompts/knowledgeBaseQA/userQuery';
+import { ChatSemanticSearchChunk } from '@/types/chunk';
+import { KnowledgeItem } from '@/types/knowledgeBase';
+
+export const knowledgeBaseQAPrompts = ({
+  chunks,
+  knowledge,
+  userQuery,
+  rewriteQuery,
+}: {
+  chunks?: ChatSemanticSearchChunk[];
+  knowledge?: KnowledgeItem[];
+  rewriteQuery?: string;
+  userQuery: string;
+}) => {
+  if ((chunks || [])?.length === 0) return '';
+
+  const domains = (knowledge || []).map((v) => v.name).join('/');
+
+  return `<knowledge_base_qa_info>
+You are also a helpful assistant good answering questions related to ${domains}. And you'll be provided with a question and several passages that might be relevant. And currently your task is to provide answer based on the question and passages.
+<knowledge_base_anwser_instruction>
+- Note that passages might not be relevant to the question, please only use the passages that are relevant.
+- if there is no relevant passage, please answer using your knowledge.
+- Answer should use the same original language as the question and follow markdown syntax.
+</knowledge_base_anwser_instruction>
+${knowledgePrompts(knowledge)}
+${chunks ? chunkPrompts(chunks) : ''}
+${userQueryPrompt(userQuery, rewriteQuery)}
+</knowledge_base_qa_info>`;
+};

package/src/prompts/knowledgeBaseQA/knowledge.ts

@@ -0,0 +1,15 @@
+import { KnowledgeItem } from '@/types/knowledgeBase';
+
+const knowledgePrompt = (item: KnowledgeItem) =>
+  `<knowledge id="${item.id}" name="${item.name}" type="${item.type}"${item.fileType ? ` fileType="${item.fileType}" ` : ''}>${item.description || ''}</knowledge>`;
+
+export const knowledgePrompts = (list?: KnowledgeItem[]) => {
+  if ((list || []).length === 0) return '';
+
+  const prompt = `<knowledge_bases>
+<knowledge_bases_docstring>here are the knowledge base scope we retrieve chunks from:</knowledge_bases_docstring>
+${list?.map((item) => knowledgePrompt(item)).join('\n')}
+</knowledge_bases>`;
+
+  return prompt.trim();
+};

package/src/prompts/knowledgeBaseQA/userQuery.ts

@@ -0,0 +1,8 @@
+export const userQueryPrompt = (userQuery: string, rewriteQuery?: string) => {
+  return `<user_query>
+<user_query_docstring>to make result better, we may rewrite user's question.If there is a rewrite query, it will be wrapper with \`rewrite_query\` tag.</user_query_docstring>
+
+<raw_query>${userQuery.trim()}</raw_query>
+${rewriteQuery ? `<rewrite_query>${rewriteQuery.trim()}</rewrite_query>` : ''}
+<user_query>`;
+};

package/src/services/chat.ts

@@ -422,7 +422,7 @@ class ChatService {
 
       const filesContext = filesPrompts({ fileList: m.fileList, imageList });
       return [
-        { text: m.content + '\n\n' + filesContext, type: 'text' },
+        { text: (m.content + '\n\n' + filesContext).trim(), type: 'text' },
         ...imageList.map(
           (i) => ({ image_url: { detail: 'auto', url: i.url }, type: 'image_url' }) as const,
         ),

package/src/store/chat/slices/aiChat/{action.test.ts → actions/__tests__/generateAIChat.test.ts}

@@ -14,7 +14,7 @@ import { sessionMetaSelectors } from '@/store/session/selectors';
 import { UploadFileItem } from '@/types/files/upload';
 import { ChatMessage } from '@/types/message';
 
-import { useChatStore } from '../../store';
+import { useChatStore } from '../../../../store';
 
 vi.stubGlobal(
   'fetch',

package/src/store/chat/slices/aiChat/actions/__tests__/rag.test.ts

@@ -0,0 +1,261 @@
+import { act, renderHook } from '@testing-library/react';
+import { Mock, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { chatService } from '@/services/chat';
+import { ragService } from '@/services/rag';
+import { useAgentStore } from '@/store/agent';
+import { agentSelectors } from '@/store/agent/selectors';
+import { chatSelectors } from '@/store/chat/selectors';
+import { systemAgentSelectors } from '@/store/user/selectors';
+import { ChatMessage } from '@/types/message';
+import { QueryRewriteSystemAgent } from '@/types/user/settings';
+
+import { useChatStore } from '../../../../store';
+
+// Mock services
+vi.mock('@/services/chat', () => ({
+  chatService: {
+    fetchPresetTaskResult: vi.fn(),
+  },
+}));
+
+vi.mock('@/services/rag', () => ({
+  ragService: {
+    deleteMessageRagQuery: vi.fn(),
+    semanticSearchForChat: vi.fn(),
+  },
+}));
+
+beforeEach(() => {
+  vi.clearAllMocks();
+});
+
+describe('chatRAG actions', () => {
+  describe('deleteUserMessageRagQuery', () => {
+    it('should not delete if message not found', async () => {
+      const { result } = renderHook(() => useChatStore());
+
+      await act(async () => {
+        await result.current.deleteUserMessageRagQuery('non-existent-id');
+      });
+
+      expect(ragService.deleteMessageRagQuery).not.toHaveBeenCalled();
+    });
+
+    it('should not delete if message has no ragQueryId', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const messageId = 'message-id';
+
+      act(() => {
+        useChatStore.setState({
+          messagesMap: {
+            default: [{ id: messageId }] as ChatMessage[],
+          },
+        });
+      });
+
+      await act(async () => {
+        await result.current.deleteUserMessageRagQuery(messageId);
+      });
+
+      expect(ragService.deleteMessageRagQuery).not.toHaveBeenCalled();
+    });
+  });
+
+  describe('internal_retrieveChunks', () => {
+    it('should retrieve chunks with existing ragQuery', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const messageId = 'message-id';
+      const existingRagQuery = 'existing-query';
+      const userQuery = 'user-query';
+
+      // Mock the message with existing ragQuery
+      vi.spyOn(chatSelectors, 'getMessageById').mockReturnValue(
+        () =>
+          ({
+            id: messageId,
+            ragQuery: existingRagQuery,
+          }) as ChatMessage,
+      );
+
+      // Mock the semantic search response
+      (ragService.semanticSearchForChat as Mock).mockResolvedValue({
+        chunks: [{ id: 'chunk-1' }],
+        queryId: 'query-id',
+      });
+
+      vi.spyOn(agentSelectors, 'currentKnowledgeIds').mockReturnValue({
+        fileIds: [],
+        knowledgeBaseIds: [],
+      });
+
+      const result1 = await act(async () => {
+        return await result.current.internal_retrieveChunks(messageId, userQuery, []);
+      });
+
+      expect(result1).toEqual({
+        chunks: [{ id: 'chunk-1' }],
+        queryId: 'query-id',
+        rewriteQuery: existingRagQuery,
+      });
+      expect(ragService.semanticSearchForChat).toHaveBeenCalledWith(
+        expect.objectContaining({
+          rewriteQuery: existingRagQuery,
+          userQuery,
+        }),
+      );
+    });
+
+    it('should rewrite query if no existing ragQuery', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const messageId = 'message-id';
+      const userQuery = 'user-query';
+      const rewrittenQuery = 'rewritten-query';
+
+      // Mock the message without ragQuery
+      vi.spyOn(chatSelectors, 'getMessageById').mockReturnValue(
+        () =>
+          ({
+            id: messageId,
+          }) as ChatMessage,
+      );
+
+      // Mock the rewrite query function
+      vi.spyOn(result.current, 'internal_rewriteQuery').mockResolvedValueOnce(rewrittenQuery);
+
+      // Mock the semantic search response
+      (ragService.semanticSearchForChat as Mock).mockResolvedValue({
+        chunks: [{ id: 'chunk-1' }],
+        queryId: 'query-id',
+      });
+
+      vi.spyOn(agentSelectors, 'currentKnowledgeIds').mockReturnValue({
+        fileIds: [],
+        knowledgeBaseIds: [],
+      });
+
+      const result2 = await act(async () => {
+        return await result.current.internal_retrieveChunks(messageId, userQuery, ['message']);
+      });
+
+      expect(result2).toEqual({
+        chunks: [{ id: 'chunk-1' }],
+        queryId: 'query-id',
+        rewriteQuery: rewrittenQuery,
+      });
+      expect(result.current.internal_rewriteQuery).toHaveBeenCalledWith(messageId, userQuery, [
+        'message',
+      ]);
+    });
+  });
+
+  describe('internal_rewriteQuery', () => {
+    it('should return original content if query rewrite is disabled', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const content = 'original content';
+
+      vi.spyOn(systemAgentSelectors, 'queryRewrite').mockReturnValueOnce({
+        enabled: false,
+      } as QueryRewriteSystemAgent);
+
+      const rewrittenQuery = await result.current.internal_rewriteQuery('id', content, []);
+
+      expect(rewrittenQuery).toBe(content);
+      expect(chatService.fetchPresetTaskResult).not.toHaveBeenCalled();
+    });
+
+    it('should rewrite query if enabled', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const messageId = 'message-id';
+      const content = 'original content';
+      const rewrittenContent = 'rewritten content';
+
+      vi.spyOn(systemAgentSelectors, 'queryRewrite').mockReturnValueOnce({
+        enabled: true,
+        model: 'gpt-3.5',
+        provider: 'openai',
+      });
+
+      (chatService.fetchPresetTaskResult as Mock).mockImplementation(({ onFinish }) => {
+        onFinish(rewrittenContent);
+      });
+
+      const rewrittenQuery = await result.current.internal_rewriteQuery(messageId, content, []);
+
+      expect(rewrittenQuery).toBe(rewrittenContent);
+      expect(chatService.fetchPresetTaskResult).toHaveBeenCalled();
+    });
+  });
+
+  describe('internal_shouldUseRAG', () => {
+    it('should return true if has enabled knowledge', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      vi.spyOn(agentSelectors, 'hasEnabledKnowledge').mockReturnValue(true);
+      vi.spyOn(chatSelectors, 'currentUserFiles').mockReturnValue([]);
+
+      expect(result.current.internal_shouldUseRAG()).toBe(true);
+    });
+
+    it('should return true if has user files', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      vi.spyOn(agentSelectors, 'hasEnabledKnowledge').mockReturnValue(false);
+      vi.spyOn(chatSelectors, 'currentUserFiles').mockReturnValue([{ id: 'file-1' }] as any);
+
+      expect(result.current.internal_shouldUseRAG()).toBe(true);
+    });
+
+    it('should return false if no knowledge or files', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      vi.spyOn(agentSelectors, 'hasEnabledKnowledge').mockReturnValue(false);
+      vi.spyOn(chatSelectors, 'currentUserFiles').mockReturnValue([]);
+
+      expect(result.current.internal_shouldUseRAG()).toBe(false);
+    });
+  });
+
+  describe('rewriteQuery', () => {
+    it('should not rewrite if message not found', async () => {
+      const { result } = renderHook(() => useChatStore());
+
+      vi.spyOn(chatSelectors, 'getMessageById').mockReturnValue(() => undefined);
+      const rewriteSpy = vi.spyOn(result.current, 'internal_rewriteQuery');
+
+      await act(async () => {
+        await result.current.rewriteQuery('non-existent-id');
+      });
+
+      expect(rewriteSpy).not.toHaveBeenCalled();
+    });
+
+    it('should rewrite query for existing message', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const messageId = 'message-id';
+      const content = 'message content';
+
+      vi.spyOn(chatSelectors, 'getMessageById').mockReturnValue(
+        () =>
+          ({
+            id: messageId,
+            content,
+          }) as ChatMessage,
+      );
+
+      vi.spyOn(chatSelectors, 'currentChatsWithHistoryConfig').mockReturnValue([
+        { content: 'history' },
+      ] as ChatMessage[]);
+
+      const rewriteSpy = vi.spyOn(result.current, 'internal_rewriteQuery');
+      const deleteSpy = vi.spyOn(result.current, 'deleteUserMessageRagQuery');
+
+      await act(async () => {
+        await result.current.rewriteQuery(messageId);
+      });
+
+      expect(deleteSpy).toHaveBeenCalledWith(messageId);
+      expect(rewriteSpy).toHaveBeenCalledWith(messageId, content, ['history']);
+    });
+  });
+});

package/src/store/chat/slices/aiChat/{action.ts → actions/generateAIChat.ts}

@@ -4,10 +4,10 @@ import { produce } from 'immer';
 import { template } from 'lodash-es';
 import { StateCreator } from 'zustand/vanilla';
 
-import { chainAnswerWithContext } from '@/chains/answerWithContext';
 import { LOADING_FLAT, MESSAGE_CANCEL_FLAT } from '@/const/message';
 import { TraceEventType, TraceNameMap } from '@/const/trace';
 import { isServerMode } from '@/const/version';
+import { knowledgeBaseQAPrompts } from '@/prompts/knowledgeBaseQA';
 import { chatService } from '@/services/chat';
 import { messageService } from '@/services/message';
 import { useAgentStore } from '@/store/agent';
@@ -16,27 +16,14 @@ import { chatHelpers } from '@/store/chat/helpers';
 import { ChatStore } from '@/store/chat/store';
 import { messageMapKey } from '@/store/chat/utils/messageMapKey';
 import { useSessionStore } from '@/store/session';
-import { UploadFileItem } from '@/types/files/upload';
-import { ChatMessage, CreateMessageParams } from '@/types/message';
+import { ChatMessage, CreateMessageParams, SendMessageParams } from '@/types/message';
 import { MessageSemanticSearchChunk } from '@/types/rag';
 import { setNamespace } from '@/utils/storeDebug';
 
-import { chatSelectors, topicSelectors } from '../../selectors';
-import { ChatRAGAction, chatRag } from './actions/rag';
+import { chatSelectors, topicSelectors } from '../../../selectors';
 
 const n = setNamespace('ai');
 
-export interface SendMessageParams {
-  message: string;
-  files?: UploadFileItem[];
-  onlyAddUserMessage?: boolean;
-  /**
-   *
-   * https://github.com/lobehub/lobe-chat/pull/2086
-   */
-  isWelcomeQuestion?: boolean;
-}
-
 interface ProcessMessageParams {
   traceId?: string;
   isWelcomeQuestion?: boolean;
@@ -46,7 +33,7 @@ interface ProcessMessageParams {
   ragQuery?: string;
 }
 
-export interface ChatAIChatAction extends ChatRAGAction {
+export interface AIGenerateAction {
   /**
    * Sends a new message to the AI chat system
    */
@@ -110,14 +97,12 @@ const getAgentConfig = () => agentSelectors.currentAgentConfig(useAgentStore.get
 const getAgentChatConfig = () => agentSelectors.currentAgentChatConfig(useAgentStore.getState());
 const getAgentKnowledge = () => agentSelectors.currentEnabledKnowledge(useAgentStore.getState());
 
-export const chatAiChat: StateCreator<
+export const generateAIChat: StateCreator<
   ChatStore,
   [['zustand/devtools', never]],
   [],
-  ChatAIChatAction
-> = (set, get, ...rest) => ({
-  ...chatRag(set, get, ...rest),
-
+  AIGenerateAction
+> = (set, get) => ({
   delAndRegenerateMessage: async (id) => {
     const traceId = chatSelectors.getTraceIdByMessageId(id)(get());
     get().internal_resendMessage(id, traceId);
@@ -284,10 +269,11 @@ export const chatAiChat: StateCreator<
 
     let fileChunks: MessageSemanticSearchChunk[] | undefined;
     let ragQueryId;
+
     // go into RAG flow if there is ragQuery flag
     if (params?.ragQuery) {
       // 1. get the relative chunks from semantic search
-      const { chunks, queryId } = await get().internal_retrieveChunks(
+      const { chunks, queryId, rewriteQuery } = await get().internal_retrieveChunks(
         userMessageId,
         params?.ragQuery,
         // should skip the last content
@@ -296,19 +282,21 @@ export const chatAiChat: StateCreator<
 
       ragQueryId = queryId;
 
+      const lastMsg = messages.pop() as ChatMessage;
+
       // 2. build the retrieve context messages
-      const retrieveContext = chainAnswerWithContext({
-        context: chunks.map((c) => c.text as string),
-        question: params?.ragQuery,
-        knowledge: getAgentKnowledge().map((knowledge) => knowledge.name),
+      const knowledgeBaseQAContext = knowledgeBaseQAPrompts({
+        chunks,
+        userQuery: lastMsg.content,
+        rewriteQuery,
+        knowledge: getAgentKnowledge(),
       });
 
       // 3. add the retrieve context messages to the messages history
-      if (retrieveContext.messages && retrieveContext.messages?.length > 0) {
-        // remove the last message due to the query is in the retrieveContext
-        messages.pop();
-        retrieveContext.messages?.forEach((m) => messages.push(m as ChatMessage));
-      }
+      messages.push({
+        ...lastMsg,
+        content: (lastMsg.content + '\n\n' + knowledgeBaseQAContext).trim(),
+      });
 
       fileChunks = chunks.map((c) => ({ id: c.id, similarity: c.similarity }));
     }
@@ -514,7 +502,7 @@ export const chatAiChat: StateCreator<
 
     await internal_coreProcessMessage(contextMessages, latestMsg.id, {
       traceId,
-      ragQuery: get().internal_shouldUseRAG() ? currentMessage.content : undefined,
+      ragQuery: get().internal_shouldUseRAG() ? latestMsg.content : undefined,
     });
   },
 

package/src/store/chat/slices/aiChat/actions/index.ts

@@ -0,0 +1,20 @@
+import { StateCreator } from 'zustand/vanilla';
+
+import { ChatStore } from '@/store/chat/store';
+
+import { AIGenerateAction, generateAIChat } from './generateAIChat';
+import { ChatRAGAction, chatRag } from './rag';
+
+export interface ChatAIChatAction extends ChatRAGAction, AIGenerateAction {
+  /**/
+}
+
+export const chatAiChat: StateCreator<
+  ChatStore,
+  [['zustand/devtools', never]],
+  [],
+  ChatAIChatAction
+> = (...params) => ({
+  ...chatRag(...params),
+  ...generateAIChat(...params),
+});

package/src/store/chat/slices/aiChat/actions/rag.ts

@@ -21,7 +21,7 @@ export interface ChatRAGAction {
     id: string,
     userQuery: string,
     messages: string[],
-  ) => Promise<{ chunks: ChatSemanticSearchChunk[]; queryId: string }>;
+  ) => Promise<{ chunks: ChatSemanticSearchChunk[]; queryId: string; rewriteQuery?: string }>;
   /**
    * Rewrite user content to better RAG query
    */
@@ -64,12 +64,11 @@ export const chatRag: StateCreator<ChatStore, [['zustand/devtools', never]], [],
     const message = chatSelectors.getMessageById(id)(get());
 
     // 1. get the rewrite query
-    let rewriteQuery = message?.ragQuery || userQuery;
+    let rewriteQuery = message?.ragQuery as string | undefined;
 
-    // only rewrite query length is less than 15 characters, refs: https://github.com/lobehub/lobe-chat/pull/4288
     // if there is no ragQuery and there is a chat history
     // we need to rewrite the user message to get better results
-    if (rewriteQuery.length < 15 && !message?.ragQuery && messages.length > 0) {
+    if (!message?.ragQuery && messages.length > 0) {
       rewriteQuery = await get().internal_rewriteQuery(id, userQuery, messages);
     }
 
@@ -79,13 +78,13 @@ export const chatRag: StateCreator<ChatStore, [['zustand/devtools', never]], [],
       fileIds: knowledgeIds().fileIds.concat(files),
       knowledgeIds: knowledgeIds().knowledgeBaseIds,
       messageId: id,
-      rewriteQuery,
+      rewriteQuery: rewriteQuery || userQuery,
       userQuery,
     });
 
     get().internal_toggleMessageRAGLoading(false, id);
 
-    return { chunks, queryId };
+    return { chunks, queryId, rewriteQuery };
   },
   internal_rewriteQuery: async (id, content, messages) => {
     let rewriteQuery = content;

package/src/store/chat/store.ts

@@ -13,7 +13,7 @@ import { ChatMessageAction, chatMessage } from './slices/message/action';
 import { ChatPluginAction, chatPlugin } from './slices/plugin/action';
 import { ShareAction, chatShare } from './slices/share/action';
 import { ChatTopicAction, chatTopic } from './slices/topic/action';
-import { chatAiChat, ChatAIChatAction } from './slices/aiChat/action';
+import { chatAiChat, ChatAIChatAction } from './slices/aiChat/actions';
 
 export interface ChatStoreAction
   extends ChatMessageAction,

package/src/types/chunk/index.ts

@@ -39,6 +39,8 @@ export interface FileChunk {
 }
 
 export interface SemanticSearchChunk {
+  fileId: string | null;
+  fileName: string | null;
   id: string;
   metadata: ChunkMetadata | null;
   pageNumber?: number | null;

package/src/types/message/index.ts

@@ -2,6 +2,7 @@ import { IPluginErrorType } from '@lobehub/chat-plugin-sdk';
 
 import { ILobeAgentRuntimeErrorType } from '@/libs/agent-runtime';
 import { ErrorType } from '@/types/fetch';
+import { UploadFileItem } from '@/types/files';
 import { MessageSemanticSearchChunk } from '@/types/rag';
 
 import { BaseDataModel } from '../meta';
@@ -131,3 +132,14 @@ export interface CreateMessageParams
   topicId?: string;
   traceId?: string;
 }
+
+export interface SendMessageParams {
+  files?: UploadFileItem[];
+  /**
+   *
+   * https://github.com/lobehub/lobe-chat/pull/2086
+   */
+  isWelcomeQuestion?: boolean;
+  message: string;
+  onlyAddUserMessage?: boolean;
+}