@lobehub/chat 1.11.9 → 1.12.0

package/src/database/server/migrations/meta/_journal.json

@@ -35,6 +35,20 @@
       "when": 1721724512422,
       "tag": "0004_add_next_auth",
       "breakpoints": true
+    },
+    {
+      "idx": 5,
+      "version": "7",
+      "when": 1722944166657,
+      "tag": "0005_pgvector",
+      "breakpoints": true
+    },
+    {
+      "idx": 6,
+      "version": "7",
+      "when": 1724089032064,
+      "tag": "0006_add_knowledge_base",
+      "breakpoints": true
     }
   ],
   "version": "6"

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

@@ -0,0 +1,155 @@
+// @vitest-environment node
+import { eq } from 'drizzle-orm';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { getTestDBInstance } from '@/database/server/core/dbForTest';
+
+import { sessionGroups, users } from '../../schemas/lobechat';
+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);
+
+beforeEach(async () => {
+  await serverDB.delete(users);
+  await serverDB.insert(users).values([{ id: userId }, { id: 'user2' }]);
+});
+
+afterEach(async () => {
+  await serverDB.delete(users).where(eq(users.id, userId));
+  await serverDB.delete(sessionGroups).where(eq(sessionGroups.userId, userId));
+});
+
+describe('SessionGroupModel', () => {
+  describe('create', () => {
+    it('should create a new session group', async () => {
+      const params = {
+        name: 'Test Group',
+        sort: 1,
+      };
+
+      const result = await sessionGroupModel.create(params);
+      expect(result.id).toBeDefined();
+      expect(result).toMatchObject({ ...params, userId });
+
+      const group = await serverDB.query.sessionGroups.findFirst({
+        where: eq(sessionGroups.id, result.id),
+      });
+      expect(group).toMatchObject({ ...params, userId });
+    });
+  });
+  describe('delete', () => {
+    it('should delete a session group by id', async () => {
+      const { id } = await sessionGroupModel.create({ name: 'Test Group' });
+
+      await sessionGroupModel.delete(id);
+
+      const group = await serverDB.query.sessionGroups.findFirst({
+        where: eq(sessionGroups.id, id),
+      });
+      expect(group).toBeUndefined();
+    });
+  });
+  describe('deleteAll', () => {
+    it('should delete all session groups for the user', async () => {
+      await sessionGroupModel.create({ name: 'Test Group 1' });
+      await sessionGroupModel.create({ name: 'Test Group 2' });
+
+      await sessionGroupModel.deleteAll();
+
+      const userGroups = await serverDB.query.sessionGroups.findMany({
+        where: eq(sessionGroups.userId, userId),
+      });
+      expect(userGroups).toHaveLength(0);
+    });
+    it('should only delete session groups for the user, not others', async () => {
+      await sessionGroupModel.create({ name: 'Test Group 1' });
+      await sessionGroupModel.create({ name: 'Test Group 333' });
+
+      const anotherSessionGroupModel = new SessionGroupModel('user2');
+      await anotherSessionGroupModel.create({ name: 'Test Group 2' });
+
+      await sessionGroupModel.deleteAll();
+
+      const userGroups = await serverDB.query.sessionGroups.findMany({
+        where: eq(sessionGroups.userId, userId),
+      });
+      const total = await serverDB.query.sessionGroups.findMany();
+      expect(userGroups).toHaveLength(0);
+      expect(total).toHaveLength(1);
+    });
+  });
+
+  describe('query', () => {
+    it('should query session groups for the user', async () => {
+      await sessionGroupModel.create({ name: 'Test Group 1', sort: 2 });
+      await sessionGroupModel.create({ name: 'Test Group 2', sort: 1 });
+
+      const userGroups = await sessionGroupModel.query();
+      expect(userGroups).toHaveLength(2);
+      expect(userGroups[0].name).toBe('Test Group 2');
+      expect(userGroups[1].name).toBe('Test Group 1');
+    });
+  });
+
+  describe('findById', () => {
+    it('should find a session group by id', async () => {
+      const { id } = await sessionGroupModel.create({ name: 'Test Group' });
+
+      const group = await sessionGroupModel.findById(id);
+      expect(group).toMatchObject({
+        id,
+        name: 'Test Group',
+        userId,
+      });
+    });
+  });
+
+  describe('update', () => {
+    it('should update a session group', async () => {
+      const { id } = await sessionGroupModel.create({ name: 'Test Group' });
+
+      await sessionGroupModel.update(id, { name: 'Updated Test Group', sort: 3 });
+
+      const updatedGroup = await serverDB.query.sessionGroups.findFirst({
+        where: eq(sessionGroups.id, id),
+      });
+      expect(updatedGroup).toMatchObject({
+        id,
+        name: 'Updated Test Group',
+        sort: 3,
+        userId,
+      });
+    });
+  });
+
+  describe('updateOrder', () => {
+    it('should update order of session groups', async () => {
+      const group1 = await sessionGroupModel.create({ name: 'Test Group 1', sort: 1 });
+      const group2 = await sessionGroupModel.create({ name: 'Test Group 2', sort: 2 });
+
+      await sessionGroupModel.updateOrder([
+        { id: group1.id, sort: 3 },
+        { id: group2.id, sort: 4 },
+      ]);
+
+      const updatedGroup1 = await serverDB.query.sessionGroups.findFirst({
+        where: eq(sessionGroups.id, group1.id),
+      });
+      const updatedGroup2 = await serverDB.query.sessionGroups.findFirst({
+        where: eq(sessionGroups.id, group2.id),
+      });
+
+      expect(updatedGroup1?.sort).toBe(3);
+      expect(updatedGroup2?.sort).toBe(4);
+    });
+  });
+});

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

@@ -0,0 +1,226 @@
+// @vitest-environment node
+import { eq } from 'drizzle-orm';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { getTestDBInstance } from '@/database/server/core/dbForTest';
+
+import {
+  agents,
+  agentsFiles,
+  agentsKnowledgeBases,
+  agentsToSessions,
+  files,
+  knowledgeBases,
+  sessions,
+  users,
+} from '../../schemas/lobechat';
+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 knowledgeBase = { id: 'kb1', userId, name: 'knowledgeBase' };
+const fileList = [
+  {
+    id: '1',
+    name: 'document.pdf',
+    url: 'https://a.com/document.pdf',
+    size: 1000,
+    fileType: 'application/pdf',
+    userId,
+  },
+  {
+    id: '2',
+    name: 'image.jpg',
+    url: 'https://a.com/image.jpg',
+    size: 500,
+    fileType: 'image/jpeg',
+    userId,
+  },
+];
+
+beforeEach(async () => {
+  await serverDB.delete(users);
+  await serverDB.insert(users).values([{ id: userId }]);
+  await serverDB.insert(knowledgeBases).values(knowledgeBase);
+  await serverDB.insert(files).values(fileList);
+});
+
+afterEach(async () => {
+  await serverDB.delete(users).where(eq(users.id, userId));
+});
+
+describe('AgentModel', () => {
+  describe('getAgentConfigById', () => {
+    it('should return agent config with assigned knowledge', async () => {
+      const agentId = 'test-agent-id';
+      await serverDB.insert(agents).values({ id: agentId, userId });
+      await serverDB
+        .insert(agentsKnowledgeBases)
+        .values({ agentId, knowledgeBaseId: 'kb1', userId });
+      await serverDB.insert(agentsFiles).values({ agentId, fileId: '1', userId });
+
+      const result = await agentModel.getAgentConfigById(agentId);
+
+      expect(result).toBeDefined();
+      expect(result.id).toBe(agentId);
+      expect(result.knowledgeBases).toHaveLength(1);
+      expect(result.files).toHaveLength(1);
+    });
+  });
+
+  describe('findBySessionId', () => {
+    it('should find agent by session id', async () => {
+      const agentId = 'test-agent-id';
+      const sessionId = 'test-session-id';
+      await serverDB.insert(agents).values({ id: agentId, userId });
+      await serverDB.insert(sessions).values({ id: sessionId, userId });
+      await serverDB.insert(agentsToSessions).values({ agentId, sessionId });
+
+      const result = await agentModel.findBySessionId(sessionId);
+
+      expect(result).toBeDefined();
+      expect(result?.id).toBe(agentId);
+    });
+  });
+
+  describe('createAgentKnowledgeBase', () => {
+    it('should create a new agent knowledge base association', async () => {
+      const agent = await serverDB
+        .insert(agents)
+        .values({ userId })
+        .returning()
+        .then((res) => res[0]);
+
+      await agentModel.createAgentKnowledgeBase(agent.id, knowledgeBase.id);
+
+      const result = await serverDB.query.agentsKnowledgeBases.findFirst({
+        where: eq(agentsKnowledgeBases.agentId, agent.id),
+      });
+
+      expect(result).toMatchObject({
+        agentId: agent.id,
+        knowledgeBaseId: knowledgeBase.id,
+        userId,
+        enabled: true,
+      });
+    });
+  });
+
+  describe('deleteAgentKnowledgeBase', () => {
+    it('should delete an agent knowledge base association', async () => {
+      const agent = await serverDB
+        .insert(agents)
+        .values({ userId })
+        .returning()
+        .then((res) => res[0]);
+      await serverDB
+        .insert(agentsKnowledgeBases)
+        .values({ agentId: agent.id, knowledgeBaseId: knowledgeBase.id, userId });
+
+      await agentModel.deleteAgentKnowledgeBase(agent.id, knowledgeBase.id);
+
+      const result = await serverDB.query.agentsKnowledgeBases.findFirst({
+        where: eq(agentsKnowledgeBases.agentId, agent.id),
+      });
+
+      expect(result).toBeUndefined();
+    });
+  });
+
+  describe('toggleKnowledgeBase', () => {
+    it('should toggle the enabled status of an agent knowledge base association', async () => {
+      const agent = await serverDB
+        .insert(agents)
+        .values({ userId })
+        .returning()
+        .then((res) => res[0]);
+
+      await serverDB
+        .insert(agentsKnowledgeBases)
+        .values({ agentId: agent.id, knowledgeBaseId: knowledgeBase.id, userId, enabled: true });
+
+      await agentModel.toggleKnowledgeBase(agent.id, knowledgeBase.id, false);
+
+      const result = await serverDB.query.agentsKnowledgeBases.findFirst({
+        where: eq(agentsKnowledgeBases.agentId, agent.id),
+      });
+
+      expect(result?.enabled).toBe(false);
+    });
+  });
+
+  describe('createAgentFiles', () => {
+    it('should create new agent file associations', async () => {
+      const agent = await serverDB
+        .insert(agents)
+        .values({ userId })
+        .returning()
+        .then((res) => res[0]);
+
+      await agentModel.createAgentFiles(agent.id, ['1', '2']);
+
+      const results = await serverDB.query.agentsFiles.findMany({
+        where: eq(agentsFiles.agentId, agent.id),
+      });
+
+      expect(results).toHaveLength(2);
+      expect(results).toEqual(
+        expect.arrayContaining([
+          expect.objectContaining({ agentId: agent.id, fileId: '1', userId, enabled: true }),
+          expect.objectContaining({ agentId: agent.id, fileId: '2', userId, enabled: true }),
+        ]),
+      );
+    });
+  });
+
+  describe('deleteAgentFile', () => {
+    it('should delete an agent file association', async () => {
+      const agent = await serverDB
+        .insert(agents)
+        .values({ userId })
+        .returning()
+        .then((res) => res[0]);
+
+      await serverDB.insert(agentsFiles).values({ agentId: agent.id, fileId: '1', userId });
+
+      await agentModel.deleteAgentFile(agent.id, '1');
+
+      const result = await serverDB.query.agentsFiles.findFirst({
+        where: eq(agentsFiles.agentId, agent.id),
+      });
+
+      expect(result).toBeUndefined();
+    });
+  });
+
+  describe('toggleFile', () => {
+    it('should toggle the enabled status of an agent file association', async () => {
+      const agent = await serverDB
+        .insert(agents)
+        .values({ userId })
+        .returning()
+        .then((res) => res[0]);
+
+      await serverDB
+        .insert(agentsFiles)
+        .values({ agentId: agent.id, fileId: '1', userId, enabled: true });
+
+      await agentModel.toggleFile(agent.id, '1', false);
+
+      const result = await serverDB.query.agentsFiles.findFirst({
+        where: eq(agentsFiles.agentId, agent.id),
+      });
+
+      expect(result?.enabled).toBe(false);
+    });
+  });
+});

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

@@ -0,0 +1,176 @@
+// @vitest-environment node
+import { eq } from 'drizzle-orm';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { getTestDBInstance } from '@/database/server/core/dbForTest';
+import { AsyncTaskStatus, AsyncTaskType } from '@/types/asyncTask';
+
+import { asyncTasks, users } from '../../schemas/lobechat';
+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);
+
+beforeEach(async () => {
+  await serverDB.delete(users);
+  await serverDB.insert(users).values([{ id: userId }]);
+});
+
+afterEach(async () => {
+  await serverDB.delete(users).where(eq(users.id, userId));
+});
+
+describe('AsyncTaskModel', () => {
+  describe('create', () => {
+    it('should create a new async task', async () => {
+      const params = {
+        type: AsyncTaskType.Chunking,
+        status: AsyncTaskStatus.Processing,
+      };
+
+      const taskId = await asyncTaskModel.create(params);
+
+      const task = await serverDB.query.asyncTasks.findFirst({
+        where: eq(asyncTasks.id, taskId),
+      });
+      expect(task).toMatchObject({ ...params, userId });
+    });
+  });
+
+  describe('delete', () => {
+    it('should delete an async task by id', async () => {
+      const { id } = await serverDB
+        .insert(asyncTasks)
+        .values({
+          type: AsyncTaskType.Chunking,
+          status: AsyncTaskStatus.Processing,
+          userId,
+        })
+        .returning()
+        .then((res) => res[0]);
+
+      await asyncTaskModel.delete(id);
+
+      const task = await serverDB.query.asyncTasks.findFirst({
+        where: eq(asyncTasks.id, id),
+      });
+      expect(task).toBeUndefined();
+    });
+  });
+
+  describe('findById', () => {
+    it('should find an async task by id', async () => {
+      const { id } = await serverDB
+        .insert(asyncTasks)
+        .values({
+          type: AsyncTaskType.Chunking,
+          status: AsyncTaskStatus.Processing,
+          userId,
+        })
+        .returning()
+        .then((res) => res[0]);
+
+      const task = await asyncTaskModel.findById(id);
+      expect(task).toBeDefined();
+      expect(task?.id).toBe(id);
+    });
+  });
+
+  describe('update', () => {
+    it('should update an async task', async () => {
+      const { id } = await serverDB
+        .insert(asyncTasks)
+        .values({
+          type: AsyncTaskType.Chunking,
+          status: AsyncTaskStatus.Processing,
+          userId,
+        })
+        .returning()
+        .then((res) => res[0]);
+
+      await asyncTaskModel.update(id, { status: AsyncTaskStatus.Success });
+
+      const updatedTask = await serverDB.query.asyncTasks.findFirst({
+        where: eq(asyncTasks.id, id),
+      });
+      expect(updatedTask?.status).toBe(AsyncTaskStatus.Success);
+    });
+  });
+
+  describe('findByIds', () => {
+    it('should find async tasks by ids and type', async () => {
+      const tasks = await serverDB
+        .insert(asyncTasks)
+        .values([
+          { type: AsyncTaskType.Chunking, status: AsyncTaskStatus.Processing, userId },
+          { type: AsyncTaskType.Chunking, status: AsyncTaskStatus.Success, userId },
+          { type: AsyncTaskType.Embedding, status: AsyncTaskStatus.Processing, userId },
+        ])
+        .returning();
+
+      const chunkTasks = await asyncTaskModel.findByIds(
+        tasks.map((t) => t.id),
+        AsyncTaskType.Chunking,
+      );
+
+      expect(chunkTasks).toHaveLength(2);
+      expect(chunkTasks.every((t) => t.type === AsyncTaskType.Chunking)).toBe(true);
+    });
+  });
+
+  describe('checkTimeoutTasks', () => {
+    it('should mark tasks as error if they timeout', async () => {
+      vi.useFakeTimers();
+
+      const { id } = await serverDB
+        .insert(asyncTasks)
+        .values({
+          type: AsyncTaskType.Chunking,
+          status: AsyncTaskStatus.Processing,
+          userId,
+          createdAt: new Date(Date.now() - ASYNC_TASK_TIMEOUT - 1000), // Make sure it's older than the timeout
+        })
+        .returning()
+        .then((res) => res[0]);
+
+      await asyncTaskModel.checkTimeoutTasks([id]);
+
+      const updatedTask = await serverDB.query.asyncTasks.findFirst({
+        where: eq(asyncTasks.id, id),
+      });
+      expect(updatedTask?.status).toBe(AsyncTaskStatus.Error);
+      expect(updatedTask?.error).toBeDefined();
+
+      vi.useRealTimers();
+    });
+
+    it('should not mark tasks as error if they are not timed out', async () => {
+      const { id } = await serverDB
+        .insert(asyncTasks)
+        .values({
+          type: AsyncTaskType.Chunking,
+          status: AsyncTaskStatus.Processing,
+          userId,
+          createdAt: new Date(), // Current time, should not timeout
+        })
+        .returning()
+        .then((res) => res[0]);
+
+      await asyncTaskModel.checkTimeoutTasks([id]);
+
+      const updatedTask = await serverDB.query.asyncTasks.findFirst({
+        where: eq(asyncTasks.id, id),
+      });
+      expect(updatedTask?.status).toBe(AsyncTaskStatus.Processing);
+      expect(updatedTask?.error).toBeNull();
+    });
+  });
+});