@lobehub/chat 0.162.25 → 0.163.0

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

@@ -0,0 +1,290 @@
+import { Column, asc, count, inArray, like, sql } from 'drizzle-orm';
+import { and, desc, eq, isNull, not, or } from 'drizzle-orm/expressions';
+
+import { appEnv } from '@/config/app';
+import { INBOX_SESSION_ID } from '@/const/session';
+import { DEFAULT_AGENT_CONFIG } from '@/const/settings';
+import { serverDB } from '@/database/server/core/db';
+import { parseAgentConfig } from '@/server/globalConfig/parseDefaultAgent';
+import { ChatSessionList, LobeAgentSession } from '@/types/session';
+import { merge } from '@/utils/merge';
+
+import {
+  AgentItem,
+  NewAgent,
+  NewSession,
+  SessionItem,
+  agents,
+  agentsToSessions,
+  sessionGroups,
+  sessions,
+} from '../schemas/lobechat';
+import { idGenerator } from '../utils/idGenerator';
+
+export class SessionModel {
+  private userId: string;
+
+  constructor(userId: string) {
+    this.userId = userId;
+  }
+  // **************** Query *************** //
+
+  async query({ current = 0, pageSize = 9999 } = {}) {
+    const offset = current * pageSize;
+
+    return serverDB.query.sessions.findMany({
+      limit: pageSize,
+      offset,
+      orderBy: [desc(sessions.updatedAt)],
+      where: and(eq(sessions.userId, this.userId), not(eq(sessions.slug, INBOX_SESSION_ID))),
+      with: { agentsToSessions: { columns: {}, with: { agent: true } }, group: true },
+    });
+  }
+
+  async queryWithGroups(): Promise<ChatSessionList> {
+    // 查询所有会话
+    const result = await this.query();
+
+    const groups = await serverDB.query.sessionGroups.findMany({
+      orderBy: [asc(sessionGroups.sort), desc(sessionGroups.createdAt)],
+      where: eq(sessions.userId, this.userId),
+    });
+
+    return {
+      sessionGroups: groups as unknown as ChatSessionList['sessionGroups'],
+      sessions: result.map((item) => this.mapSessionItem(item as any)),
+    };
+  }
+
+  async queryByKeyword(keyword: string) {
+    if (!keyword) return [];
+
+    const keywordLowerCase = keyword.toLowerCase();
+
+    const data = await this.findSessions({ keyword: keywordLowerCase });
+
+    return data.map((item) => this.mapSessionItem(item as any));
+  }
+
+  async findByIdOrSlug(
+    idOrSlug: string,
+  ): Promise<(SessionItem & { agent: AgentItem }) | undefined> {
+    const result = await serverDB.query.sessions.findFirst({
+      where: and(
+        or(eq(sessions.id, idOrSlug), eq(sessions.slug, idOrSlug)),
+        eq(sessions.userId, this.userId),
+      ),
+      with: { agentsToSessions: { columns: {}, with: { agent: true } }, group: true },
+    });
+
+    if (!result) return;
+
+    return { ...result, agent: (result?.agentsToSessions?.[0] as any)?.agent } as any;
+  }
+
+  async count() {
+    const result = await serverDB
+      .select({
+        count: count(),
+      })
+      .from(sessions)
+      .where(eq(sessions.userId, this.userId))
+      .execute();
+
+    return result[0].count;
+  }
+
+  // **************** Create *************** //
+
+  async create({
+    id = idGenerator('sessions'),
+    type = 'agent',
+    session = {},
+    config = {},
+    slug,
+  }: {
+    config?: Partial<NewAgent>;
+    id?: string;
+    session?: Partial<NewSession>;
+    slug?: string;
+    type: 'agent' | 'group';
+  }): Promise<SessionItem> {
+    return serverDB.transaction(async (trx) => {
+      const newAgents = await trx
+        .insert(agents)
+        .values({
+          ...config,
+          createdAt: new Date(),
+          id: idGenerator('agents'),
+          updatedAt: new Date(),
+          userId: this.userId,
+        })
+        .returning();
+
+      const result = await trx
+        .insert(sessions)
+        .values({
+          ...session,
+          createdAt: new Date(),
+          id,
+          slug,
+          type,
+          updatedAt: new Date(),
+          userId: this.userId,
+        })
+        .returning();
+
+      await trx.insert(agentsToSessions).values({
+        agentId: newAgents[0].id,
+        sessionId: id,
+      });
+
+      return result[0];
+    });
+  }
+
+  async createInbox() {
+    const serverAgentConfig = parseAgentConfig(appEnv.DEFAULT_AGENT_CONFIG) || {};
+
+    return await this.create({
+      config: merge(DEFAULT_AGENT_CONFIG, serverAgentConfig),
+      slug: INBOX_SESSION_ID,
+      type: 'agent',
+    });
+  }
+
+  async batchCreate(newSessions: NewSession[]) {
+    const sessionsToInsert = newSessions.map((s) => {
+      return {
+        ...s,
+        id: this.genId(),
+        userId: this.userId,
+      };
+    });
+
+    return serverDB.insert(sessions).values(sessionsToInsert);
+  }
+
+  async duplicate(id: string, newTitle?: string) {
+    const result = await this.findByIdOrSlug(id);
+
+    if (!result) return;
+
+    const { agent, ...session } = result;
+    const sessionId = this.genId();
+
+    return this.create({
+      config: agent,
+      id: sessionId,
+      session: {
+        ...session,
+        title: newTitle || session.title,
+      },
+      type: 'agent',
+    });
+  }
+
+  // **************** Delete *************** //
+
+  /**
+   * Delete a session, also delete all messages and topics associated with it.
+   */
+  async delete(id: string) {
+    return serverDB
+      .delete(sessions)
+      .where(and(eq(sessions.id, id), eq(sessions.userId, this.userId)));
+  }
+
+  /**
+   * Batch delete sessions, also delete all messages and topics associated with them.
+   */
+  async batchDelete(ids: string[]) {
+    return serverDB
+      .delete(sessions)
+      .where(and(inArray(sessions.id, ids), eq(sessions.userId, this.userId)));
+  }
+
+  async deleteAll() {
+    return serverDB.delete(sessions).where(eq(sessions.userId, this.userId));
+  }
+  // **************** Update *************** //
+
+  async update(id: string, data: Partial<SessionItem>) {
+    return serverDB
+      .update(sessions)
+      .set(data)
+      .where(and(eq(sessions.id, id), eq(sessions.userId, this.userId)))
+      .returning();
+  }
+
+  async updateConfig(id: string, data: Partial<AgentItem>) {
+    return serverDB
+      .update(agents)
+      .set(data)
+      .where(and(eq(agents.id, id), eq(agents.userId, this.userId)));
+  }
+
+  // **************** Helper *************** //
+
+  private genId = () => idGenerator('sessions');
+
+  private mapSessionItem = ({
+    agentsToSessions,
+    title,
+    backgroundColor,
+    description,
+    avatar,
+    groupId,
+    ...res
+  }: SessionItem & { agentsToSessions?: { agent: AgentItem }[] }): LobeAgentSession => {
+    // TODO: 未来这里需要更好的实现方案，目前只取第一个
+    const agent = agentsToSessions?.[0]?.agent;
+    return {
+      ...res,
+      group: groupId,
+      meta: {
+        avatar: agent?.avatar ?? avatar ?? undefined,
+        backgroundColor: agent?.backgroundColor ?? backgroundColor ?? undefined,
+        description: agent?.description ?? description ?? undefined,
+        title: agent?.title ?? title ?? undefined,
+      },
+      model: agent?.model,
+    } as any;
+  };
+
+  async findSessions(params: {
+    current?: number;
+    group?: string;
+    keyword?: string;
+    pageSize?: number;
+    pinned?: boolean;
+  }) {
+    const { pinned, keyword, group, pageSize = 9999, current = 0 } = params;
+
+    const offset = current * pageSize;
+    return serverDB.query.sessions.findMany({
+      limit: pageSize,
+      offset,
+      orderBy: [desc(sessions.updatedAt)],
+      where: and(
+        eq(sessions.userId, this.userId),
+        pinned !== undefined ? eq(sessions.pinned, pinned) : eq(sessions.userId, this.userId),
+        keyword
+          ? or(
+              like(
+                sql`lower(${sessions.title})` as unknown as Column,
+                `%${keyword.toLowerCase()}%`,
+              ),
+              like(
+                sql`lower(${sessions.description})` as unknown as Column,
+                `%${keyword.toLowerCase()}%`,
+              ),
+            )
+          : eq(sessions.userId, this.userId),
+        group ? eq(sessions.groupId, group) : isNull(sessions.groupId),
+      ),
+
+      with: { agentsToSessions: { columns: {}, with: { agent: true } }, group: true },
+    });
+  }
+}

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

@@ -0,0 +1,69 @@
+import { eq } from 'drizzle-orm';
+import { and, asc, desc } from 'drizzle-orm/expressions';
+
+import { serverDB } from '@/database/server';
+import { idGenerator } from '@/database/server/utils/idGenerator';
+
+import { SessionGroupItem, sessionGroups } from '../schemas/lobechat';
+
+export class SessionGroupModel {
+  private userId: string;
+
+  constructor(userId: string) {
+    this.userId = userId;
+  }
+
+  create = async (params: { name: string; sort?: number }) => {
+    const [result] = await serverDB
+      .insert(sessionGroups)
+      .values({ ...params, id: this.genId(), userId: this.userId })
+      .returning();
+
+    return result;
+  };
+
+  delete = async (id: string) => {
+    return serverDB
+      .delete(sessionGroups)
+      .where(and(eq(sessionGroups.id, id), eq(sessionGroups.userId, this.userId)));
+  };
+
+  deleteAll = async () => {
+    return serverDB.delete(sessionGroups);
+  };
+
+  query = async () => {
+    return serverDB.query.sessionGroups.findMany({
+      orderBy: [asc(sessionGroups.sort), desc(sessionGroups.createdAt)],
+      where: eq(sessionGroups.userId, this.userId),
+    });
+  };
+
+  findById = async (id: string) => {
+    return serverDB.query.sessionGroups.findFirst({
+      where: and(eq(sessionGroups.id, id), eq(sessionGroups.userId, this.userId)),
+    });
+  };
+
+  async update(id: string, value: Partial<SessionGroupItem>) {
+    return serverDB
+      .update(sessionGroups)
+      .set({ ...value, updatedAt: new Date() })
+      .where(and(eq(sessionGroups.id, id), eq(sessionGroups.userId, this.userId)));
+  }
+
+  async updateOrder(sortMap: { id: string; sort: number }[]) {
+    await serverDB.transaction(async (tx) => {
+      const updates = sortMap.map(({ id, sort }) => {
+        return tx
+          .update(sessionGroups)
+          .set({ sort, updatedAt: new Date() })
+          .where(and(eq(sessionGroups.id, id), eq(sessionGroups.userId, this.userId)));
+      });
+
+      await Promise.all(updates);
+    });
+  }
+
+  private genId = () => idGenerator('sessionGroups');
+}

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

@@ -0,0 +1,265 @@
+import { Column, count, inArray, sql } from 'drizzle-orm';
+import { and, desc, eq, exists, isNull, like, or } from 'drizzle-orm/expressions';
+
+import { serverDB } from '@/database/server/core/db';
+
+import { NewMessage, TopicItem, messages, topics } from '../schemas/lobechat';
+import { idGenerator } from '../utils/idGenerator';
+
+export interface CreateTopicParams {
+  favorite?: boolean;
+  messages?: string[];
+  sessionId?: string | null;
+  title: string;
+}
+
+interface QueryTopicParams {
+  current?: number;
+  pageSize?: number;
+  sessionId?: string | null;
+}
+
+export class TopicModel {
+  private userId: string;
+
+  constructor(userId: string) {
+    this.userId = userId;
+  }
+  // **************** Query *************** //
+
+  async query({ current = 0, pageSize = 9999, sessionId }: QueryTopicParams = {}) {
+    const offset = current * pageSize;
+
+    return (
+      serverDB
+        .select({
+          createdAt: topics.createdAt,
+          favorite: topics.favorite,
+          id: topics.id,
+          title: topics.title,
+          updatedAt: topics.updatedAt,
+        })
+        .from(topics)
+        .where(and(eq(topics.userId, this.userId), this.matchSession(sessionId)))
+        // In boolean sorting, false is considered "smaller" than true.
+        // So here we use desc to ensure that topics with favorite as true are in front.
+        .orderBy(desc(topics.favorite), desc(topics.updatedAt))
+        .limit(pageSize)
+        .offset(offset)
+    );
+  }
+
+  async findById(id: string) {
+    return serverDB.query.topics.findFirst({
+      where: and(eq(topics.id, id), eq(topics.userId, this.userId)),
+    });
+  }
+
+  async queryAll(): Promise<TopicItem[]> {
+    return serverDB
+      .select()
+      .from(topics)
+      .orderBy(topics.updatedAt)
+      .where(eq(topics.userId, this.userId))
+      .execute();
+  }
+
+  async queryByKeyword(keyword: string, sessionId?: string | null): Promise<TopicItem[]> {
+    if (!keyword) return [];
+
+    const keywordLowerCase = keyword.toLowerCase();
+
+    const matchKeyword = (field: any) =>
+      like(sql`lower(${field})` as unknown as Column, `%${keywordLowerCase}%`);
+
+    return serverDB.query.topics.findMany({
+      orderBy: [desc(topics.updatedAt)],
+      where: and(
+        eq(topics.userId, this.userId),
+        this.matchSession(sessionId),
+        or(
+          matchKeyword(topics.title),
+          exists(
+            serverDB
+              .select()
+              .from(messages)
+              .where(and(eq(messages.topicId, topics.id), or(matchKeyword(messages.content)))),
+          ),
+        ),
+      ),
+    });
+  }
+
+  async count() {
+    const result = await serverDB
+      .select({
+        count: count(),
+      })
+      .from(topics)
+      .where(eq(topics.userId, this.userId))
+      .execute();
+
+    return result[0].count;
+  }
+
+  // **************** Create *************** //
+
+  async create(
+    { messages: messageIds, ...params }: CreateTopicParams,
+    id: string = this.genId(),
+  ): Promise<TopicItem> {
+    return serverDB.transaction(async (tx) => {
+      // 在 topics 表中插入新的 topic
+      const [topic] = await tx
+        .insert(topics)
+        .values({
+          ...params,
+          id: id,
+          userId: this.userId,
+        })
+        .returning();
+
+      // 如果有关联的 messages, 更新它们的 topicId
+      if (messageIds && messageIds.length > 0) {
+        await tx
+          .update(messages)
+          .set({ topicId: topic.id })
+          .where(and(eq(messages.userId, this.userId), inArray(messages.id, messageIds)));
+      }
+
+      return topic;
+    });
+  }
+
+  async batchCreate(topicParams: (CreateTopicParams & { id?: string })[]) {
+    // 开始一个事务
+    return serverDB.transaction(async (tx) => {
+      // 在 topics 表中批量插入新的 topics
+      const createdTopics = await tx
+        .insert(topics)
+        .values(
+          topicParams.map((params) => ({
+            favorite: params.favorite,
+            id: params.id || this.genId(),
+            sessionId: params.sessionId,
+            title: params.title,
+            userId: this.userId,
+          })),
+        )
+        .returning();
+
+      // 对每个新创建的 topic,更新关联的 messages 的 topicId
+      await Promise.all(
+        createdTopics.map(async (topic, index) => {
+          const messageIds = topicParams[index].messages;
+          if (messageIds && messageIds.length > 0) {
+            await tx
+              .update(messages)
+              .set({ topicId: topic.id })
+              .where(and(eq(messages.userId, this.userId), inArray(messages.id, messageIds)));
+          }
+        }),
+      );
+
+      return createdTopics;
+    });
+  }
+
+  async duplicate(topicId: string, newTitle?: string) {
+    return serverDB.transaction(async (tx) => {
+      // find original topic
+      const originalTopic = await tx.query.topics.findFirst({
+        where: and(eq(topics.id, topicId), eq(topics.userId, this.userId)),
+      });
+
+      if (!originalTopic) {
+        throw new Error(`Topic with id ${topicId} not found`);
+      }
+
+      // copy topic
+      const [duplicatedTopic] = await tx
+        .insert(topics)
+        .values({
+          ...originalTopic,
+          id: this.genId(),
+          title: newTitle || originalTopic?.title,
+        })
+        .returning();
+
+      // 查找与原始 topic 关联的 messages
+      const originalMessages = await tx
+        .select()
+        .from(messages)
+        .where(and(eq(messages.topicId, topicId), eq(messages.userId, this.userId)));
+
+      // copy messages
+      const duplicatedMessages = await Promise.all(
+        originalMessages.map(async (message) => {
+          const result = (await tx
+            .insert(messages)
+            .values({
+              ...message,
+              id: idGenerator('messages'),
+              topicId: duplicatedTopic.id,
+            })
+            .returning()) as NewMessage[];
+
+          return result[0];
+        }),
+      );
+
+      return {
+        messages: duplicatedMessages,
+        topic: duplicatedTopic,
+      };
+    });
+  }
+
+  // **************** Delete *************** //
+
+  /**
+   * Delete a session, also delete all messages and topics associated with it.
+   */
+  async delete(id: string) {
+    return serverDB.delete(topics).where(and(eq(topics.id, id), eq(topics.userId, this.userId)));
+  }
+
+  /**
+   * Deletes multiple topics based on the sessionId.
+   */
+  async batchDeleteBySessionId(sessionId?: string | null) {
+    return serverDB
+      .delete(topics)
+      .where(and(this.matchSession(sessionId), eq(topics.userId, this.userId)));
+  }
+
+  /**
+   * Deletes multiple topics and all messages associated with them in a transaction.
+   */
+  async batchDelete(ids: string[]) {
+    return serverDB
+      .delete(topics)
+      .where(and(inArray(topics.id, ids), eq(topics.userId, this.userId)));
+  }
+
+  async deleteAll() {
+    return serverDB.delete(topics).where(eq(topics.userId, this.userId));
+  }
+
+  // **************** Update *************** //
+
+  async update(id: string, data: Partial<TopicItem>) {
+    return serverDB
+      .update(topics)
+      .set({ ...data, updatedAt: new Date() })
+      .where(and(eq(topics.id, id), eq(topics.userId, this.userId)))
+      .returning();
+  }
+
+  // **************** Helper *************** //
+
+  private genId = () => idGenerator('topics');
+
+  private matchSession = (sessionId?: string | null) =>
+    sessionId ? eq(topics.sessionId, sessionId) : isNull(topics.sessionId);
+}