lobe-chat

1
import { BaseModel } from '@/database/client/core';
2
import { DBModel } from '@/database/client/core/types/db';
3
import { MessageModel } from '@/database/client/models/message';
4
import { DB_Topic, DB_TopicSchema } from '@/database/client/schemas/topic';
5
import { ChatTopic } from '@/types/topic';
6
import { nanoid } from '@/utils/uuid';
7

8
export interface CreateTopicParams {
9
  favorite?: boolean;
10
  messages?: string[];
11
  sessionId: string;
12
  title: string;
13
}
14

15
export interface QueryTopicParams {
16
  current?: number;
17
  pageSize?: number;
18
  sessionId: string;
19
}
20

21
class _TopicModel extends BaseModel {
22
  constructor() {
23
    super('topics', DB_TopicSchema);
24
  }
25

26
  // **************** Query *************** //
27

28
  async query({ pageSize = 9999, current = 0, sessionId }: QueryTopicParams): Promise<ChatTopic[]> {
29
    const offset = current * pageSize;
30

31
    // get all topics
32
    const allTopics = await this.table.where('sessionId').equals(sessionId).toArray();
33

34
    // 将所有主题按星标消息优先，时间倒序进行排序
35
    const sortedTopics = allTopics.sort((a, b) => {
36
      if (a.favorite && !b.favorite) return -1; // a是星标，b不是，a排前面
37
      if (!a.favorite && b.favorite) return 1; // b是星标，a不是，b排前面
38

39
      // 如果星标状态相同，则按时间倒序排序
40
      return b.createdAt - a.createdAt;
41
    });
42

43
    // handle pageSize
44
    const pagedTopics = sortedTopics.slice(offset, offset + pageSize);
45

46
    return pagedTopics.map((i) => this.mapToChatTopic(i));
47
  }
48

49
  queryAll() {
50
    return this.table.orderBy('updatedAt').toArray();
51
  }
52

53
  /**
54
   * Query topics by keyword in title, message content, or translated content
55
   * @param keyword The keyword to search for
56
   * @param sessionId The currently activated session id.
57
   */
58
  async queryByKeyword(keyword: string, sessionId?: string): Promise<ChatTopic[]> {
59
    if (!keyword) return [];
60

61
    console.time('queryTopicsByKeyword');
62
    const keywordLowerCase = keyword.toLowerCase();
63

64
    // Find topics with matching title
65
    const queryTable = sessionId ? this.table.where('sessionId').equals(sessionId) : this.table;
66
    const matchingTopicsPromise = queryTable
67
      .filter((topic) => topic.title.toLowerCase().includes(keywordLowerCase))
68
      .toArray();
69

70
    // Find messages with matching content or translate.content
71
    const queryMessages = sessionId
72
      ? this.db.messages.where('sessionId').equals(sessionId)
73
      : this.db.messages;
74
    const matchingMessagesPromise = queryMessages
75
      .filter((message) => {
76
        // check content
77
        if (message.content.toLowerCase().includes(keywordLowerCase)) return true;
78

79
        // check translate content
80
        if (message.translate && message.translate.content) {
81
          return message.translate.content.toLowerCase().includes(keywordLowerCase);
82
        }
83

84
        return false;
85
      })
86
      .toArray();
87

88
    // Resolve both promises
89
    const [matchingTopics, matchingMessages] = await Promise.all([
90
      matchingTopicsPromise,
91
      matchingMessagesPromise,
92
    ]);
93

94
    // Extract topic IDs from messages
95
    const topicIdsFromMessages = matchingMessages.map((message) => message.topicId);
96

97
    // Combine topic IDs from both sources
98
    const combinedTopicIds = new Set([
99
      ...topicIdsFromMessages,
100
      ...matchingTopics.map((topic) => topic.id),
101
    ]);
102

103
    // Retrieve unique topics by IDs
104
    const uniqueTopics = await this.table
105
      .where('id')
106
      .anyOf([...combinedTopicIds])
107
      .toArray();
108

109
    console.timeEnd('queryTopicsByKeyword');
110
    return uniqueTopics.map((i) => ({ ...i, favorite: !!i.favorite }));
111
  }
112

113
  async findBySessionId(sessionId: string) {
114
    return this.table.where({ sessionId }).toArray();
115
  }
116

117
  async findById(id: string): Promise<DBModel<DB_Topic>> {
118
    return this.table.get(id);
119
  }
120

121
  async count() {
122
    return this.table.count();
123
  }
124

125
  // **************** Create *************** //
126

127
  async create({ title, favorite, sessionId, messages }: CreateTopicParams, id = nanoid()) {
128
    const topic = await this._addWithSync(
129
      { favorite: favorite ? 1 : 0, sessionId, title: title },
130
      id,
131
    );
132

133
    // add topicId to these messages
134
    if (messages) {
135
      await MessageModel.batchUpdate(messages, { topicId: topic.id });
136
    }
137

138
    return topic;
139
  }
140

141
  async batchCreate(topics: CreateTopicParams[]) {
142
    return this._batchAdd(topics.map((t) => ({ ...t, favorite: t.favorite ? 1 : 0 })));
143
  }
144

145
  async duplicateTopic(topicId: string, newTitle?: string) {
146
    return this.db.transaction('rw', [this.db.topics, this.db.messages], async () => {
147
      // Step 1: get DB_Topic
148
      const topic = await this.findById(topicId);
149

150
      if (!topic) {
151
        throw new Error(`Topic with id ${topicId} not found`);
152
      }
153

154
      // Step 3: 查询与 `topic` 关联的 `messages`
155
      const originalMessages = await MessageModel.queryByTopicId(topicId);
156

157
      const duplicateMessages = await MessageModel.duplicateMessages(originalMessages);
158

159
      const { id } = await this.create({
160
        ...this.mapToChatTopic(topic),
161
        messages: duplicateMessages.map((m) => m.id),
162
        sessionId: topic.sessionId!,
163
        title: newTitle || topic.title,
164
      });
165

166
      return id;
167
    });
168
  }
169

170
  // **************** Delete *************** //
171

172
  /**
173
   * Deletes a topic and all messages associated with it.
174
   */
175
  async delete(id: string) {
176
    return this.db.transaction('rw', [this.table, this.db.messages], async () => {
177
      // Delete all messages associated with the topic
178
      await MessageModel.batchDeleteByTopicId(id);
179

180
      await this._deleteWithSync(id);
181
    });
182
  }
183

184
  /**
185
   * Deletes multiple topic based on the sessionId.
186
   *
187
   * @param {string} sessionId - The identifier of the assistant associated with the messages.
188
   * @returns {Promise<void>}
189
   */
190
  async batchDeleteBySessionId(sessionId: string): Promise<void> {
191
    // use sessionId as the filter criteria in the query.
192
    const query = this.table.where('sessionId').equals(sessionId);
193

194
    // Retrieve a collection of message IDs that satisfy the criteria
195
    const topicIds = await query.primaryKeys();
196

197
    // Use the bulkDelete method to delete all selected messages in bulk
198
    return this._bulkDeleteWithSync(topicIds);
199
  }
200
  /**
201
   * Deletes multiple topics and all messages associated with them in a transaction.
202
   */
203
  async batchDelete(topicIds: string[]) {
204
    return this.db.transaction('rw', [this.table, this.db.messages], async () => {
205
      // Iterate over each topicId and delete related messages, then delete the topic itself
206
      for (const topicId of topicIds) {
207
        // Delete all messages associated with the topic
208
        await this.delete(topicId);
209
      }
210
    });
211
  }
212

213
  async clearTable() {
214
    return this._clearWithSync();
215
  }
216

217
  // **************** Update *************** //
218
  async update(id: string, data: Partial<DB_Topic>) {
219
    return super._updateWithSync(id, data);
220
  }
221

222
  async toggleFavorite(id: string, newState?: boolean) {
223
    const topic = await this.findById(id);
224
    if (!topic) {
225
      throw new Error(`Topic with id ${id} not found`);
226
    }
227

228
    // Toggle the 'favorite' status
229
    const nextState = typeof newState !== 'undefined' ? newState : !topic.favorite;
230

231
    await this.update(id, { favorite: nextState ? 1 : 0 });
232

233
    return nextState;
234
  }
235

236
  // **************** Helper *************** //
237

238
  private mapToChatTopic = (dbTopic: DBModel<DB_Topic>): ChatTopic => ({
239
    ...dbTopic,
240
    favorite: !!dbTopic.favorite,
241
  });
242
}
243

244
export const TopicModel = new _TopicModel();
245