@lobehub/chat 1.11.9 → 1.12.1

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

@@ -0,0 +1,132 @@
+// @vitest-environment node
+import { eq } from 'drizzle-orm';
+import { desc } from 'drizzle-orm/expressions';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { getTestDBInstance } from '@/database/server/core/dbForTest';
+
+import { NewKnowledgeBase, knowledgeBases, users } from '../../schemas/lobechat';
+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);
+
+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(knowledgeBases).where(eq(knowledgeBases.userId, userId));
+});
+
+describe('KnowledgeBaseModel', () => {
+  describe('create', () => {
+    it('should create a new knowledge base', async () => {
+      const params = {
+        name: 'Test Group',
+      } as NewKnowledgeBase;
+
+      const result = await knowledgeBaseModel.create(params);
+      expect(result.id).toBeDefined();
+      expect(result).toMatchObject({ ...params, userId });
+
+      const group = await serverDB.query.knowledgeBases.findFirst({
+        where: eq(knowledgeBases.id, result.id),
+      });
+      expect(group).toMatchObject({ ...params, userId });
+    });
+  });
+  describe('delete', () => {
+    it('should delete a knowledge base by id', async () => {
+      const { id } = await knowledgeBaseModel.create({ name: 'Test Group' });
+
+      await knowledgeBaseModel.delete(id);
+
+      const group = await serverDB.query.knowledgeBases.findFirst({
+        where: eq(knowledgeBases.id, id),
+      });
+      expect(group).toBeUndefined();
+    });
+  });
+  describe('deleteAll', () => {
+    it('should delete all knowledge bases for the user', async () => {
+      await knowledgeBaseModel.create({ name: 'Test Group 1' });
+      await knowledgeBaseModel.create({ name: 'Test Group 2' });
+
+      await knowledgeBaseModel.deleteAll();
+
+      const userGroups = await serverDB.query.knowledgeBases.findMany({
+        where: eq(knowledgeBases.userId, userId),
+      });
+      expect(userGroups).toHaveLength(0);
+    });
+    it('should only delete knowledge bases for the user, not others', async () => {
+      await knowledgeBaseModel.create({ name: 'Test Group 1' });
+      await knowledgeBaseModel.create({ name: 'Test Group 333' });
+
+      const anotherSessionGroupModel = new KnowledgeBaseModel('user2');
+      await anotherSessionGroupModel.create({ name: 'Test Group 2' });
+
+      await knowledgeBaseModel.deleteAll();
+
+      const userGroups = await serverDB.query.knowledgeBases.findMany({
+        where: eq(knowledgeBases.userId, userId),
+      });
+      const total = await serverDB.query.knowledgeBases.findMany();
+      expect(userGroups).toHaveLength(0);
+      expect(total).toHaveLength(1);
+    });
+  });
+
+  describe('query', () => {
+    it('should query knowledge bases for the user', async () => {
+      await knowledgeBaseModel.create({ name: 'Test Group 1' });
+      await knowledgeBaseModel.create({ name: 'Test Group 2' });
+
+      const userGroups = await knowledgeBaseModel.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 knowledge base by id', async () => {
+      const { id } = await knowledgeBaseModel.create({ name: 'Test Group' });
+
+      const group = await knowledgeBaseModel.findById(id);
+      expect(group).toMatchObject({
+        id,
+        name: 'Test Group',
+        userId,
+      });
+    });
+  });
+
+  describe('update', () => {
+    it('should update a knowledge base', async () => {
+      const { id } = await knowledgeBaseModel.create({ name: 'Test Group' });
+
+      await knowledgeBaseModel.update(id, { name: 'Updated Test Group' });
+
+      const updatedGroup = await serverDB.query.knowledgeBases.findFirst({
+        where: eq(knowledgeBases.id, id),
+      });
+      expect(updatedGroup).toMatchObject({
+        id,
+        name: 'Updated Test Group',
+        userId,
+      });
+    });
+  });
+});

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

@@ -5,11 +5,11 @@ import { getTestDBInstance } from '@/database/server/core/dbForTest';
 
 import {
   files,
-  filesToMessages,
   messagePlugins,
   messageTTS,
   messageTranslates,
   messages,
+  messagesFiles,
   sessions,
   topics,
   users,
@@ -202,7 +202,7 @@ describe('MessageModel', () => {
           .insert(messageTTS)
           .values([{ id: '1' }, { id: '2', voice: 'a', fileId: 'f-1', contentMd5: 'abc' }]);
 
-        await trx.insert(filesToMessages).values([
+        await trx.insert(messagesFiles).values([
           { fileId: 'f-0', messageId: '1' },
           { fileId: 'f-3', messageId: '1' },
         ]);
@@ -214,10 +214,13 @@ describe('MessageModel', () => {
       // 断言结果
       expect(result).toHaveLength(2);
       expect(result[0].id).toBe('1');
-      expect(result[0].files).toEqual(['f-0', 'f-3']);
+      expect(result[0].imageList).toEqual([
+        { alt: 'file-1', id: 'f-0', url: expect.stringContaining('/abc') },
+        { alt: 'file-3', id: 'f-3', url: expect.stringContaining('/abc') },
+      ]);
 
       expect(result[1].id).toBe('2');
-      expect(result[1].files).toEqual([]);
+      expect(result[1].imageList).toEqual([]);
     });
 
     it('should include translate, tts and other extra fields in query result', async () => {

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

@@ -13,7 +13,12 @@ export class TemplateModel {
   }
 
   create = async (params: NewSessionGroup) => {
-    return serverDB.insert(sessionGroups).values({ ...params, userId: this.userId });
+    const [result] = await serverDB
+      .insert(sessionGroups)
+      .values({ ...params, userId: this.userId })
+      .returning();
+
+    return result;
   };
 
   delete = async (id: string) => {
@@ -22,6 +27,10 @@ export class TemplateModel {
       .where(and(eq(sessionGroups.id, id), eq(sessionGroups.userId, this.userId)));
   };
 
+  deleteAll = async () => {
+    return serverDB.delete(sessionGroups).where(eq(sessionGroups.userId, this.userId));
+  };
+
   query = async () => {
     return serverDB.query.sessionGroups.findMany({
       orderBy: [desc(sessionGroups.updatedAt)],

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

@@ -0,0 +1,165 @@
+import { inArray } from 'drizzle-orm';
+import { and, desc, eq } from 'drizzle-orm/expressions';
+
+import { serverDB } from '@/database/server';
+import {
+  agents,
+  agentsFiles,
+  agentsKnowledgeBases,
+  agentsToSessions,
+  files,
+  knowledgeBases,
+} from '@/database/server/schemas/lobechat';
+
+export class AgentModel {
+  private userId: string;
+  constructor(userId: string) {
+    this.userId = userId;
+  }
+
+  async getAgentConfigById(id: string) {
+    const agent = await serverDB.query.agents.findFirst({ where: eq(agents.id, id) });
+
+    const knowledge = await this.getAgentAssignedKnowledge(id);
+
+    return { ...agent, ...knowledge };
+  }
+
+  async getAgentAssignedKnowledge(id: string) {
+    const knowledgeBaseResult = await serverDB
+      .select({ enabled: agentsKnowledgeBases.enabled, knowledgeBases })
+      .from(agentsKnowledgeBases)
+      .where(eq(agentsKnowledgeBases.agentId, id))
+      .orderBy(desc(agentsKnowledgeBases.createdAt))
+      .leftJoin(knowledgeBases, eq(knowledgeBases.id, agentsKnowledgeBases.knowledgeBaseId));
+
+    const fileResult = await serverDB
+      .select({ enabled: agentsFiles.enabled, files })
+      .from(agentsFiles)
+      .where(eq(agentsFiles.agentId, id))
+      .orderBy(desc(agentsFiles.createdAt))
+      .leftJoin(files, eq(files.id, agentsFiles.fileId));
+
+    return {
+      files: fileResult.map((item) => ({
+        ...item.files,
+        enabled: item.enabled,
+      })),
+      knowledgeBases: knowledgeBaseResult.map((item) => ({
+        ...item.knowledgeBases,
+        enabled: item.enabled,
+      })),
+    };
+  }
+
+  /**
+   * Find agent by session id
+   */
+  async findBySessionId(sessionId: string) {
+    const item = await serverDB.query.agentsToSessions.findFirst({
+      where: eq(agentsToSessions.sessionId, sessionId),
+    });
+    if (!item) return;
+
+    const agentId = item.agentId;
+
+    return this.getAgentConfigById(agentId);
+  }
+
+  createAgentKnowledgeBase = async (
+    agentId: string,
+    knowledgeBaseId: string,
+    enabled: boolean = true,
+  ) => {
+    return serverDB
+      .insert(agentsKnowledgeBases)
+      .values({
+        agentId,
+        enabled,
+        knowledgeBaseId,
+        userId: this.userId,
+      })
+      .execute();
+  };
+
+  deleteAgentKnowledgeBase = async (agentId: string, knowledgeBaseId: string) => {
+    return serverDB
+      .delete(agentsKnowledgeBases)
+      .where(
+        and(
+          eq(agentsKnowledgeBases.agentId, agentId),
+          eq(agentsKnowledgeBases.knowledgeBaseId, knowledgeBaseId),
+          eq(agentsKnowledgeBases.userId, this.userId),
+        ),
+      )
+      .execute();
+  };
+
+  toggleKnowledgeBase = async (agentId: string, knowledgeBaseId: string, enabled?: boolean) => {
+    return serverDB
+      .update(agentsKnowledgeBases)
+      .set({ enabled })
+      .where(
+        and(
+          eq(agentsKnowledgeBases.agentId, agentId),
+          eq(agentsKnowledgeBases.knowledgeBaseId, knowledgeBaseId),
+          eq(agentsKnowledgeBases.userId, this.userId),
+        ),
+      )
+      .execute();
+  };
+
+  createAgentFiles = async (agentId: string, fileIds: string[], enabled: boolean = true) => {
+    // Exclude the fileIds that already exist in agentsFiles, and then insert them
+    const existingFiles = await serverDB
+      .select({ id: agentsFiles.fileId })
+      .from(agentsFiles)
+      .where(
+        and(
+          eq(agentsFiles.agentId, agentId),
+          eq(agentsFiles.userId, this.userId),
+          inArray(agentsFiles.fileId, fileIds),
+        ),
+      );
+
+    const existingFilesIds = new Set(existingFiles.map((item) => item.id));
+
+    const needToInsertFileIds = fileIds.filter((fileId) => !existingFilesIds.has(fileId));
+
+    if (needToInsertFileIds.length === 0) return;
+
+    return serverDB
+      .insert(agentsFiles)
+      .values(
+        needToInsertFileIds.map((fileId) => ({ agentId, enabled, fileId, userId: this.userId })),
+      )
+      .execute();
+  };
+
+  deleteAgentFile = async (agentId: string, fileId: string) => {
+    return serverDB
+      .delete(agentsFiles)
+      .where(
+        and(
+          eq(agentsFiles.agentId, agentId),
+          eq(agentsFiles.fileId, fileId),
+          eq(agentsFiles.userId, this.userId),
+        ),
+      )
+      .execute();
+  };
+
+  toggleFile = async (agentId: string, fileId: string, enabled?: boolean) => {
+    return serverDB
+      .update(agentsFiles)
+      .set({ enabled })
+      .where(
+        and(
+          eq(agentsFiles.agentId, agentId),
+          eq(agentsFiles.fileId, fileId),
+          eq(agentsFiles.userId, this.userId),
+        ),
+      )
+      .execute();
+  };
+}

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

@@ -0,0 +1,96 @@
+import { eq, inArray, lt } from 'drizzle-orm';
+import { and } from 'drizzle-orm/expressions';
+
+import { serverDB } from '@/database/server';
+import {
+  AsyncTaskError,
+  AsyncTaskErrorType,
+  AsyncTaskStatus,
+  AsyncTaskType,
+} from '@/types/asyncTask';
+
+import { AsyncTaskSelectItem, NewAsyncTaskItem, asyncTasks } from '../schemas/lobechat';
+
+// set timeout to about 5 minutes, and give 2s padding time
+export const ASYNC_TASK_TIMEOUT = 298 * 1000;
+
+export class AsyncTaskModel {
+  private userId: string;
+
+  constructor(userId: string) {
+    this.userId = userId;
+  }
+
+  create = async (params: Pick<NewAsyncTaskItem, 'type' | 'status'>): Promise<string> => {
+    const data = await serverDB
+      .insert(asyncTasks)
+      .values({ ...params, userId: this.userId })
+      .returning();
+
+    return data[0].id;
+  };
+
+  delete = async (id: string) => {
+    return serverDB
+      .delete(asyncTasks)
+      .where(and(eq(asyncTasks.id, id), eq(asyncTasks.userId, this.userId)));
+  };
+
+  findById = async (id: string) => {
+    return serverDB.query.asyncTasks.findFirst({ where: and(eq(asyncTasks.id, id)) });
+  };
+
+  update(taskId: string, value: Partial<AsyncTaskSelectItem>) {
+    return serverDB
+      .update(asyncTasks)
+      .set({ ...value, updatedAt: new Date() })
+      .where(and(eq(asyncTasks.id, taskId)));
+  }
+
+  findByIds = async (taskIds: string[], type: AsyncTaskType): Promise<AsyncTaskSelectItem[]> => {
+    let chunkTasks: AsyncTaskSelectItem[] = [];
+
+    if (taskIds.length > 0) {
+      await this.checkTimeoutTasks(taskIds);
+      chunkTasks = await serverDB.query.asyncTasks.findMany({
+        where: and(inArray(asyncTasks.id, taskIds), eq(asyncTasks.type, type)),
+      });
+    }
+
+    return chunkTasks;
+  };
+
+  /**
+   * make the task status to be `error` if the task is not finished in 20 seconds
+   */
+  async checkTimeoutTasks(ids: string[]) {
+    const tasks = await serverDB
+      .select({ id: asyncTasks.id })
+      .from(asyncTasks)
+      .where(
+        and(
+          inArray(asyncTasks.id, ids),
+          eq(asyncTasks.status, AsyncTaskStatus.Processing),
+          lt(asyncTasks.createdAt, new Date(Date.now() - ASYNC_TASK_TIMEOUT)),
+        ),
+      );
+
+    if (tasks.length > 0) {
+      await serverDB
+        .update(asyncTasks)
+        .set({
+          error: {
+            body: { detail: 'chunking task is timeout, please try again' },
+            name: AsyncTaskErrorType.Timeout,
+          } as AsyncTaskError,
+          status: AsyncTaskStatus.Error,
+        })
+        .where(
+          inArray(
+            asyncTasks.id,
+            tasks.map((item) => item.id),
+          ),
+        );
+    }
+  }
+}

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

@@ -0,0 +1,203 @@
+import { asc, cosineDistance, count, eq, inArray, sql } from 'drizzle-orm';
+import { and, desc } from 'drizzle-orm/expressions';
+
+import { serverDB } from '@/database/server';
+import { ChunkMetadata, FileChunk, SemanticSearchChunk } from '@/types/chunk';
+
+import {
+  NewChunkItem,
+  NewUnstructuredChunkItem,
+  chunks,
+  embeddings,
+  fileChunks,
+  files,
+  unstructuredChunks,
+} from '../schemas/lobechat';
+
+export class ChunkModel {
+  private userId: string;
+
+  constructor(userId: string) {
+    this.userId = userId;
+  }
+
+  bulkCreate = async (params: NewChunkItem[], fileId: string) => {
+    return serverDB.transaction(async (trx) => {
+      const result = await trx.insert(chunks).values(params).returning();
+
+      const fileChunksData = result.map((chunk) => ({ chunkId: chunk.id, fileId }));
+
+      if (fileChunksData.length > 0) {
+        await trx.insert(fileChunks).values(fileChunksData);
+      }
+
+      return result;
+    });
+  };
+
+  bulkCreateUnstructuredChunks = async (params: NewUnstructuredChunkItem[]) => {
+    return serverDB.insert(unstructuredChunks).values(params);
+  };
+
+  delete = async (id: string) => {
+    return serverDB.delete(chunks).where(and(eq(chunks.id, id), eq(chunks.userId, this.userId)));
+  };
+
+  findById = async (id: string) => {
+    return serverDB.query.chunks.findFirst({
+      where: and(eq(chunks.id, id)),
+    });
+  };
+
+  async findByFileId(id: string, page = 0) {
+    const data = await serverDB
+      .select({
+        abstract: chunks.abstract,
+        createdAt: chunks.createdAt,
+        id: chunks.id,
+        index: chunks.index,
+        metadata: chunks.metadata,
+        text: chunks.text,
+        type: chunks.type,
+        updatedAt: chunks.updatedAt,
+      })
+      .from(chunks)
+      .innerJoin(fileChunks, eq(chunks.id, fileChunks.chunkId))
+      .where(and(eq(fileChunks.fileId, id), eq(chunks.userId, this.userId)))
+      .limit(20)
+      .offset(page * 20)
+      .orderBy(asc(chunks.index));
+
+    return data.map((item) => {
+      const metadata = item.metadata as ChunkMetadata;
+
+      return { ...item, metadata, pageNumber: metadata?.pageNumber } as FileChunk;
+    });
+  }
+
+  async getChunksTextByFileId(id: string): Promise<{ id: string; text: string }[]> {
+    const data = await serverDB
+      .select()
+      .from(chunks)
+      .innerJoin(fileChunks, eq(chunks.id, fileChunks.chunkId))
+      .where(eq(fileChunks.fileId, id));
+
+    return data
+      .map((item) => item.chunks)
+      .map((chunk) => ({ id: chunk.id, text: this.mapChunkText(chunk) }))
+      .filter((chunk) => chunk.text) as { id: string; text: string }[];
+  }
+
+  async countByFileIds(ids: string[]) {
+    if (ids.length === 0) return [];
+
+    return serverDB
+      .select({
+        count: count(fileChunks.chunkId),
+        id: fileChunks.fileId,
+      })
+      .from(fileChunks)
+      .where(inArray(fileChunks.fileId, ids))
+      .groupBy(fileChunks.fileId);
+  }
+
+  async countByFileId(ids: string) {
+    const data = await serverDB
+      .select({
+        count: count(fileChunks.chunkId),
+        id: fileChunks.fileId,
+      })
+      .from(fileChunks)
+      .where(eq(fileChunks.fileId, ids))
+      .groupBy(fileChunks.fileId);
+
+    return data[0]?.count ?? 0;
+  }
+
+  async semanticSearch({
+    embedding,
+    fileIds,
+  }: {
+    embedding: number[];
+    fileIds: string[] | undefined;
+    query: string;
+  }) {
+    const similarity = sql<number>`1 - (${cosineDistance(embeddings.embeddings, embedding)})`;
+
+    const data = await serverDB
+      .select({
+        id: chunks.id,
+        index: chunks.index,
+        metadata: chunks.metadata,
+        similarity,
+        text: chunks.text,
+        type: chunks.type,
+      })
+      .from(chunks)
+      .leftJoin(embeddings, eq(chunks.id, embeddings.chunkId))
+      .leftJoin(fileChunks, eq(chunks.id, fileChunks.chunkId))
+      .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,
+      }),
+    );
+  }
+
+  async semanticSearchForChat({
+    embedding,
+    fileIds,
+  }: {
+    embedding: number[];
+    fileIds: string[] | undefined;
+    query: string;
+  }) {
+    const similarity = sql<number>`1 - (${cosineDistance(embeddings.embeddings, embedding)})`;
+
+    const result = await serverDB
+      .select({
+        fileId: files.id,
+        filename: files.name,
+        id: chunks.id,
+        index: chunks.index,
+        metadata: chunks.metadata,
+        similarity,
+        text: chunks.text,
+        type: chunks.type,
+      })
+      .from(chunks)
+      .leftJoin(embeddings, eq(chunks.id, embeddings.chunkId))
+      .leftJoin(fileChunks, eq(chunks.id, fileChunks.chunkId))
+      .leftJoin(files, eq(files.id, fileChunks.fileId))
+      .where(and(fileIds && fileIds.length > 0 ? inArray(fileChunks.fileId, fileIds) : undefined))
+      .orderBy((t) => desc(t.similarity))
+      .limit(5);
+
+    return result.map((item) => {
+      return {
+        id: item.id,
+        index: item.index,
+        similarity: item.similarity,
+        text: this.mapChunkText(item),
+      };
+    });
+  }
+
+  private mapChunkText = (chunk: { metadata: any; text: string | null; type: string | null }) => {
+    let text = chunk.text;
+
+    if (chunk.type === 'Table') {
+      text = `${chunk.text}
+
+content in Table html is below:
+${(chunk.metadata as ChunkMetadata).text_as_html}
+`;
+    }
+
+    return text;
+  };
+}

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

@@ -0,0 +1,50 @@
+import { eq } from 'drizzle-orm';
+import { and } from 'drizzle-orm/expressions';
+
+import { serverDB } from '@/database/server';
+
+import { NewEmbeddingsItem, embeddings } from '../schemas/lobechat';
+
+export class EmbeddingModel {
+  private userId: string;
+
+  constructor(userId: string) {
+    this.userId = userId;
+  }
+
+  create = async (value: Omit<NewEmbeddingsItem, 'userId'>) => {
+    const [item] = await serverDB
+      .insert(embeddings)
+      .values({ ...value, userId: this.userId })
+      .returning();
+
+    return item.id as string;
+  };
+
+  bulkCreate = async (values: Omit<NewEmbeddingsItem, 'userId'>[]) => {
+    return serverDB
+      .insert(embeddings)
+      .values(values.map((item) => ({ ...item, userId: this.userId })))
+      .onConflictDoNothing({
+        target: [embeddings.chunkId],
+      });
+  };
+
+  delete = async (id: string) => {
+    return serverDB
+      .delete(embeddings)
+      .where(and(eq(embeddings.id, id), eq(embeddings.userId, this.userId)));
+  };
+
+  query = async () => {
+    return serverDB.query.embeddings.findMany({
+      where: eq(embeddings.userId, this.userId),
+    });
+  };
+
+  findById = async (id: string) => {
+    return serverDB.query.embeddings.findFirst({
+      where: and(eq(embeddings.id, id), eq(embeddings.userId, this.userId)),
+    });
+  };
+}