@lobehub/chat 1.35.9 → 1.35.10

package/CHANGELOG.md

@@ -2,6 +2,31 @@
 
 # Changelog
 
+### [Version 1.35.10](https://github.com/lobehub/lobe-chat/compare/v1.35.9...v1.35.10)
+
+<sup>Released on **2024-12-03**</sup>
+
+#### ♻ Code Refactoring
+
+- **misc**: Refactor the server db model implement.
+
+<br/>
+
+<details>
+<summary><kbd>Improvements and Fixes</kbd></summary>
+
+#### Code refactoring
+
+- **misc**: Refactor the server db model implement, closes [#4878](https://github.com/lobehub/lobe-chat/issues/4878) ([3814853](https://github.com/lobehub/lobe-chat/commit/3814853))
+
+</details>
+
+<div align="right">
+
+[![](https://img.shields.io/badge/-BACK_TO_TOP-151515?style=flat-square)](#readme-top)
+
+</div>
+
 ### [Version 1.35.9](https://github.com/lobehub/lobe-chat/compare/v1.35.8...v1.35.9)
 
 <sup>Released on **2024-12-03**</sup>

package/changelog/v1.json

@@ -1,4 +1,13 @@
 [
+  {
+    "children": {
+      "improvements": [
+        "Refactor the server db model implement."
+      ]
+    },
+    "date": "2024-12-03",
+    "version": "1.35.10"
+  },
   {
     "children": {},
     "date": "2024-12-03",

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@lobehub/chat",
-  "version": "1.35.9",
+  "version": "1.35.10",
   "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/app/(main)/repos/[id]/@menu/default.tsx

@@ -1,6 +1,7 @@
 import { notFound } from 'next/navigation';
 import { Flexbox } from 'react-layout-kit';
 
+import { serverDB } from '@/database/server';
 import { KnowledgeBaseModel } from '@/database/server/models/knowledgeBase';
 
 import Head from './Head';
@@ -14,7 +15,7 @@ type Props = { params: Params };
 
 const MenuPage = async ({ params }: Props) => {
   const id = params.id;
-  const item = await KnowledgeBaseModel.findById(params.id);
+  const item = await KnowledgeBaseModel.findById(serverDB, params.id);
 
   if (!item) return notFound();
 

package/src/app/(main)/repos/[id]/page.tsx

@@ -1,5 +1,6 @@
 import { redirect } from 'next/navigation';
 
+import { serverDB } from '@/database/server';
 import { KnowledgeBaseModel } from '@/database/server/models/knowledgeBase';
 import FileManager from '@/features/FileManager';
 
@@ -10,7 +11,7 @@ interface Params {
 type Props = { params: Params };
 
 export default async ({ params }: Props) => {
-  const item = await KnowledgeBaseModel.findById(params.id);
+  const item = await KnowledgeBaseModel.findById(serverDB, params.id);
 
   if (!item) return redirect('/repos');
 

package/src/database/schemas/topic.ts

@@ -3,6 +3,8 @@ import { boolean, jsonb, pgTable, text, unique } from 'drizzle-orm/pg-core';
 import { createInsertSchema } from 'drizzle-zod';
 
 import { idGenerator } from '@/database/utils/idGenerator';
+import { ChatTopicMetadata } from '@/types/topic';
+
 import { timestamps, timestamptz } from './_helpers';
 import { sessions } from './session';
 import { users } from './user';
@@ -21,7 +23,7 @@ export const topics = pgTable(
       .notNull(),
     clientId: text('client_id'),
     historySummary: text('history_summary'),
-    metadata: jsonb('metadata'),
+    metadata: jsonb('metadata').$type<ChatTopicMetadata | undefined>(),
     ...timestamps,
   },
   (t) => ({

package/src/database/server/core/dbForTest.ts

@@ -11,6 +11,8 @@ import { serverDBEnv } from '@/config/db';
 
 import * as schema from '../../schemas';
 
+const migrationsFolder = join(__dirname, '../../migrations');
+
 export const getTestDBInstance = async () => {
   let connectionString = serverDBEnv.DATABASE_TEST_URL;
 
@@ -23,9 +25,7 @@ export const getTestDBInstance = async () => {
 
     const db = nodeDrizzle(client, { schema });
 
-    await nodeMigrator.migrate(db, {
-      migrationsFolder: join(__dirname, '../../migrations'),
-    });
+    await nodeMigrator.migrate(db, { migrationsFolder });
 
     return db;
   }
@@ -37,9 +37,7 @@ export const getTestDBInstance = async () => {
 
   const db = neonDrizzle(client, { schema });
 
-  await migrator.migrate(db, {
-    migrationsFolder: join(__dirname, '../migrations'),
-  });
+  await migrator.migrate(db, { migrationsFolder });
 
   return db;
 };

package/src/database/server/models/__tests__/_test_template.ts

@@ -1,6 +1,6 @@
 // @vitest-environment node
 import { eq } from 'drizzle-orm';
-import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+import { afterEach, beforeEach, describe, expect, it } from 'vitest';
 
 import { getTestDBInstance } from '@/database/server/core/dbForTest';
 
@@ -9,14 +9,8 @@ import { SessionGroupModel } from '../sessionGroup';
 
 let serverDB = await getTestDBInstance();
 
-vi.mock('@/database/server/core/db', async () => ({
-  get serverDB() {
-    return serverDB;
-  },
-}));
-
 const userId = 'session-group-model-test-user-id';
-const sessionGroupModel = new SessionGroupModel(userId);
+const sessionGroupModel = new SessionGroupModel(serverDB, userId);
 
 beforeEach(async () => {
   await serverDB.delete(users);
@@ -74,7 +68,7 @@ describe('SessionGroupModel', () => {
       await sessionGroupModel.create({ name: 'Test Group 1' });
       await sessionGroupModel.create({ name: 'Test Group 333' });
 
-      const anotherSessionGroupModel = new SessionGroupModel('user2');
+      const anotherSessionGroupModel = new SessionGroupModel(serverDB, 'user2');
       await anotherSessionGroupModel.create({ name: 'Test Group 2' });
 
       await sessionGroupModel.deleteAll();

package/src/database/server/models/__tests__/agent.test.ts

@@ -1,6 +1,6 @@
 // @vitest-environment node
 import { eq } from 'drizzle-orm';
-import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+import { afterEach, beforeEach, describe, expect, it } from 'vitest';
 
 import { getTestDBInstance } from '@/database/server/core/dbForTest';
 
@@ -18,14 +18,8 @@ import { AgentModel } from '../agent';
 
 let serverDB = await getTestDBInstance();
 
-vi.mock('@/database/server/core/db', async () => ({
-  get serverDB() {
-    return serverDB;
-  },
-}));
-
 const userId = 'agent-model-test-user-id';
-const agentModel = new AgentModel(userId);
+const agentModel = new AgentModel(serverDB, userId);
 
 const knowledgeBase = { id: 'kb1', userId, name: 'knowledgeBase' };
 const fileList = [

package/src/database/server/models/__tests__/asyncTask.test.ts

@@ -10,14 +10,8 @@ import { ASYNC_TASK_TIMEOUT, AsyncTaskModel } from '../asyncTask';
 
 let serverDB = await getTestDBInstance();
 
-vi.mock('@/database/server/core/db', async () => ({
-  get serverDB() {
-    return serverDB;
-  },
-}));
-
 const userId = 'async-task-model-test-user-id';
-const asyncTaskModel = new AsyncTaskModel(userId);
+const asyncTaskModel = new AsyncTaskModel(serverDB, userId);
 
 beforeEach(async () => {
   await serverDB.delete(users);

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

@@ -1,30 +1,18 @@
 // @vitest-environment node
 import { eq } from 'drizzle-orm';
-import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+import { afterEach, beforeEach, describe, expect, it } from 'vitest';
 
 import { getTestDBInstance } from '@/database/server/core/dbForTest';
+import { uuid } from '@/utils/uuid';
 
-import {
-  chunks,
-  embeddings,
-  fileChunks,
-  files,
-  unstructuredChunks,
-  users,
-} from '../../../schemas';
+import { chunks, embeddings, fileChunks, files, unstructuredChunks, users } from '../../../schemas';
 import { ChunkModel } from '../chunk';
 import { codeEmbedding, designThinkingQuery, designThinkingQuery2 } from './fixtures/embedding';
 
 let serverDB = await getTestDBInstance();
 
-vi.mock('@/database/server/core/db', async () => ({
-  get serverDB() {
-    return serverDB;
-  },
-}));
-
 const userId = 'chunk-model-test-user-id';
-const chunkModel = new ChunkModel(userId);
+const chunkModel = new ChunkModel(serverDB, userId);
 const sharedFileList = [
   {
     id: '1',
@@ -79,6 +67,27 @@ describe('ChunkModel', () => {
       expect(createdChunks[0]).toMatchObject(params[0]);
       expect(createdChunks[1]).toMatchObject(params[1]);
     });
+
+    // 测试空参数场景
+    it('should handle empty params array', async () => {
+      const result = await chunkModel.bulkCreate([], '1');
+      expect(result).toHaveLength(0);
+    });
+
+    // 测试事务回滚
+    it('should rollback transaction on error', async () => {
+      const invalidParams = [
+        { text: 'Chunk 1', userId },
+        { index: 'abc', userId }, // 这会导致错误
+      ] as any;
+
+      await expect(chunkModel.bulkCreate(invalidParams, '1')).rejects.toThrow();
+
+      const createdChunks = await serverDB.query.chunks.findMany({
+        where: eq(chunks.userId, userId),
+      });
+      expect(createdChunks).toHaveLength(0);
+    });
   });
 
   describe('delete', () => {
@@ -191,6 +200,41 @@ describe('ChunkModel', () => {
       expect(result[1].id).toBe(chunk2.id);
       expect(result[0].similarity).toBeGreaterThan(result[1].similarity);
     });
+    // 补充无文件 ID 的搜索场景
+    it('should perform semantic search without fileIds', async () => {
+      const [chunk1, chunk2] = await serverDB
+        .insert(chunks)
+        .values([
+          { text: 'Test Chunk 1', userId },
+          { text: 'Test Chunk 2', userId },
+        ])
+        .returning();
+
+      await serverDB.insert(embeddings).values([
+        { chunkId: chunk1.id, embeddings: designThinkingQuery, userId },
+        { chunkId: chunk2.id, embeddings: codeEmbedding, userId },
+      ]);
+
+      const result = await chunkModel.semanticSearch({
+        embedding: designThinkingQuery2,
+        fileIds: undefined,
+        query: 'design thinking',
+      });
+
+      expect(result).toBeDefined();
+      expect(result).toHaveLength(2);
+    });
+
+    // 测试空结果场景
+    it('should return empty array when no matches found', async () => {
+      const result = await chunkModel.semanticSearch({
+        embedding: designThinkingQuery,
+        fileIds: ['non-existent-file'],
+        query: 'no matches',
+      });
+
+      expect(result).toHaveLength(0);
+    });
   });
 
   describe('bulkCreateUnstructuredChunks', () => {
@@ -391,5 +435,100 @@ content in Table html is below:
 <table>...</table>
 `);
     });
+
+    it('should handle null text', () => {
+      const chunk = {
+        text: null,
+        type: 'Text',
+        metadata: {},
+      };
+
+      const result = chunkModel['mapChunkText'](chunk);
+      expect(result).toBeNull();
+    });
+
+    it('should handle missing metadata for Table type', () => {
+      const chunk = {
+        text: 'Table text',
+        type: 'Table',
+        metadata: {},
+      };
+
+      const result = chunkModel['mapChunkText'](chunk);
+      expect(result).toContain('Table text');
+      expect(result).toContain('content in Table html is below:');
+      expect(result).toContain('undefined'); // metadata.text_as_html is undefined
+    });
+  });
+
+  describe('findById', () => {
+    it('should find a chunk by id', async () => {
+      // Create a test chunk
+      const [chunk] = await serverDB
+        .insert(chunks)
+        .values({ text: 'Test Chunk', userId })
+        .returning();
+
+      const result = await chunkModel.findById(chunk.id);
+
+      expect(result).toBeDefined();
+      expect(result?.id).toBe(chunk.id);
+      expect(result?.text).toBe('Test Chunk');
+    });
+
+    it('should return null for non-existent id', async () => {
+      const result = await chunkModel.findById(uuid());
+      expect(result).toBeUndefined();
+    });
+  });
+
+  describe('semanticSearchForChat', () => {
+    // 测试空文件 ID 列表场景
+    it('should return empty array when fileIds is empty', async () => {
+      const result = await chunkModel.semanticSearchForChat({
+        embedding: designThinkingQuery,
+        fileIds: [],
+        query: 'test',
+      });
+
+      expect(result).toHaveLength(0);
+    });
+
+    // 测试结果限制
+    it('should limit results to 5 items', async () => {
+      const fileId = '1';
+      // Create 6 chunks
+      const chunkResult = await serverDB
+        .insert(chunks)
+        .values(
+          Array(6)
+            .fill(0)
+            .map((_, i) => ({ text: `Test Chunk ${i}`, userId })),
+        )
+        .returning();
+
+      await serverDB.insert(fileChunks).values(
+        chunkResult.map((chunk) => ({
+          fileId,
+          chunkId: chunk.id,
+        })),
+      );
+
+      await serverDB.insert(embeddings).values(
+        chunkResult.map((chunk) => ({
+          chunkId: chunk.id,
+          embeddings: designThinkingQuery,
+          userId,
+        })),
+      );
+
+      const result = await chunkModel.semanticSearchForChat({
+        embedding: designThinkingQuery2,
+        fileIds: [fileId],
+        query: 'test',
+      });
+
+      expect(result).toHaveLength(5);
+    });
   });
 });

package/src/database/server/models/__tests__/file.test.ts

@@ -2,27 +2,14 @@
 import { eq, inArray } from 'drizzle-orm';
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 
-import { getServerDBConfig, serverDBEnv } from '@/config/db';
 import { getTestDBInstance } from '@/database/server/core/dbForTest';
 import { FilesTabs, SortType } from '@/types/files';
 
-import {
-  files,
-  globalFiles,
-  knowledgeBaseFiles,
-  knowledgeBases,
-  users,
-} from '../../../schemas';
+import { files, globalFiles, knowledgeBaseFiles, knowledgeBases, users } from '../../../schemas';
 import { FileModel } from '../file';
 
 let serverDB = await getTestDBInstance();
 
-vi.mock('@/database/server/core/db', async () => ({
-  get serverDB() {
-    return serverDB;
-  },
-}));
-
 let DISABLE_REMOVE_GLOBAL_FILE = false;
 
 vi.mock('@/config/db', async () => ({
@@ -38,7 +25,7 @@ vi.mock('@/config/db', async () => ({
 }));
 
 const userId = 'file-model-test-user-id';
-const fileModel = new FileModel(userId);
+const fileModel = new FileModel(serverDB, userId);
 
 const knowledgeBase = { id: 'kb1', userId, name: 'knowledgeBase' };
 beforeEach(async () => {
@@ -603,4 +590,125 @@ describe('FileModel', () => {
       expect(size).toBe(3500);
     });
   });
+
+  describe('findByNames', () => {
+    it('should find files by names', async () => {
+      // 准备测试数据
+      const fileList = [
+        {
+          name: 'test1.txt',
+          url: 'https://example.com/test1.txt',
+          size: 100,
+          fileType: 'text/plain',
+          userId,
+        },
+        {
+          name: 'test2.txt',
+          url: 'https://example.com/test2.txt',
+          size: 200,
+          fileType: 'text/plain',
+          userId,
+        },
+        {
+          name: 'other.txt',
+          url: 'https://example.com/other.txt',
+          size: 300,
+          fileType: 'text/plain',
+          userId,
+        },
+      ];
+
+      await serverDB.insert(files).values(fileList);
+
+      // 测试查找文件
+      const result = await fileModel.findByNames(['test1', 'test2']);
+      expect(result).toHaveLength(2);
+      expect(result.map((f) => f.name)).toContain('test1.txt');
+      expect(result.map((f) => f.name)).toContain('test2.txt');
+    });
+
+    it('should return empty array when no files match names', async () => {
+      const result = await fileModel.findByNames(['nonexistent']);
+      expect(result).toHaveLength(0);
+    });
+
+    it('should only find files belonging to current user', async () => {
+      // 准备测试数据
+      await serverDB.insert(files).values([
+        {
+          name: 'test1.txt',
+          url: 'https://example.com/test1.txt',
+          size: 100,
+          fileType: 'text/plain',
+          userId,
+        },
+        {
+          name: 'test2.txt',
+          url: 'https://example.com/test2.txt',
+          size: 200,
+          fileType: 'text/plain',
+          userId: 'user2', // 不同用户的文件
+        },
+      ]);
+
+      const result = await fileModel.findByNames(['test']);
+      expect(result).toHaveLength(1);
+      expect(result[0].name).toBe('test1.txt');
+    });
+  });
+
+  describe('deleteGlobalFile', () => {
+    it('should delete global file by hashId', async () => {
+      // 准备测试数据
+      const globalFile = {
+        hashId: 'test-hash',
+        fileType: 'text/plain',
+        size: 100,
+        url: 'https://example.com/global-file.txt',
+        metadata: { key: 'value' },
+      };
+
+      await serverDB.insert(globalFiles).values(globalFile);
+
+      // 执行删除操作
+      await fileModel.deleteGlobalFile('test-hash');
+
+      // 验证文件已被删除
+      const result = await serverDB.query.globalFiles.findFirst({
+        where: eq(globalFiles.hashId, 'test-hash'),
+      });
+      expect(result).toBeUndefined();
+    });
+
+    it('should not throw error when deleting non-existent global file', async () => {
+      // 删除不存在的文件不应抛出错误
+      await expect(fileModel.deleteGlobalFile('non-existent-hash')).resolves.not.toThrow();
+    });
+
+    it('should only delete specified global file', async () => {
+      // 准备测试数据
+      const globalFiles1 = {
+        hashId: 'hash1',
+        fileType: 'text/plain',
+        size: 100,
+        url: 'https://example.com/file1.txt',
+      };
+      const globalFiles2 = {
+        hashId: 'hash2',
+        fileType: 'text/plain',
+        size: 200,
+        url: 'https://example.com/file2.txt',
+      };
+
+      await serverDB.insert(globalFiles).values([globalFiles1, globalFiles2]);
+
+      // 删除一个文件
+      await fileModel.deleteGlobalFile('hash1');
+
+      // 验证只有指定文件被删除
+      const remainingFiles = await serverDB.query.globalFiles.findMany();
+      expect(remainingFiles).toHaveLength(1);
+      expect(remainingFiles[0].hashId).toBe('hash2');
+    });
+  });
 });

package/src/database/server/models/__tests__/knowledgeBase.test.ts

@@ -1,7 +1,7 @@
 // @vitest-environment node
 import { eq } from 'drizzle-orm';
 import { and, desc } from 'drizzle-orm/expressions';
-import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+import { afterEach, beforeEach, describe, expect, it } from 'vitest';
 
 import { getTestDBInstance } from '@/database/server/core/dbForTest';
 
@@ -17,14 +17,8 @@ import { KnowledgeBaseModel } from '../knowledgeBase';
 
 let serverDB = await getTestDBInstance();
 
-vi.mock('@/database/server/core/db', async () => ({
-  get serverDB() {
-    return serverDB;
-  },
-}));
-
 const userId = 'session-group-model-test-user-id';
-const knowledgeBaseModel = new KnowledgeBaseModel(userId);
+const knowledgeBaseModel = new KnowledgeBaseModel(serverDB, userId);
 
 beforeEach(async () => {
   await serverDB.delete(users);
@@ -82,7 +76,7 @@ describe('KnowledgeBaseModel', () => {
       await knowledgeBaseModel.create({ name: 'Test Group 1' });
       await knowledgeBaseModel.create({ name: 'Test Group 333' });
 
-      const anotherSessionGroupModel = new KnowledgeBaseModel('user2');
+      const anotherSessionGroupModel = new KnowledgeBaseModel(serverDB, 'user2');
       await anotherSessionGroupModel.create({ name: 'Test Group 2' });
 
       await knowledgeBaseModel.deleteAll();
@@ -235,7 +229,7 @@ describe('KnowledgeBaseModel', () => {
     it('should find a knowledge base by id without user restriction', async () => {
       const { id } = await knowledgeBaseModel.create({ name: 'Test Group' });
 
-      const group = await KnowledgeBaseModel.findById(id);
+      const group = await KnowledgeBaseModel.findById(serverDB, id);
       expect(group).toMatchObject({
         id,
         name: 'Test Group',
@@ -244,10 +238,10 @@ describe('KnowledgeBaseModel', () => {
     });
 
     it('should find a knowledge base created by another user', async () => {
-      const anotherKnowledgeBaseModel = new KnowledgeBaseModel('user2');
+      const anotherKnowledgeBaseModel = new KnowledgeBaseModel(serverDB, 'user2');
       const { id } = await anotherKnowledgeBaseModel.create({ name: 'Another User Group' });
 
-      const group = await KnowledgeBaseModel.findById(id);
+      const group = await KnowledgeBaseModel.findById(serverDB, id);
       expect(group).toMatchObject({
         id,
         name: 'Another User Group',