lobe-chat

1
import Dexie, { Transaction } from 'dexie';
2

3
import { MigrationLLMSettings } from '@/migrations/FromV3ToV4';
4
import { MigrationAgentChatConfig } from '@/migrations/FromV5ToV6';
5
import { MigrationKeyValueSettings } from '@/migrations/FromV6ToV7';
6
import { uuid } from '@/utils/uuid';
7

8
import { DB_File } from '../schemas/files';
9
import { DB_Message } from '../schemas/message';
10
import { DB_Plugin } from '../schemas/plugin';
11
import { DB_Session } from '../schemas/session';
12
import { DB_SessionGroup } from '../schemas/sessionGroup';
13
import { DB_Topic } from '../schemas/topic';
14
import { DB_User } from '../schemas/user';
15
import { migrateSettingsToUser } from './migrations/migrateSettingsToUser';
16
import {
17
  dbSchemaV1,
18
  dbSchemaV2,
19
  dbSchemaV3,
20
  dbSchemaV4,
21
  dbSchemaV5,
22
  dbSchemaV6,
23
  dbSchemaV7,
24
  dbSchemaV9,
25
} from './schemas';
26
import { DBModel, LOBE_CHAT_LOCAL_DB_NAME } from './types/db';
27

28
export interface LobeDBSchemaMap {
29
  files: DB_File;
30
  messages: DB_Message;
31
  plugins: DB_Plugin;
32
  sessionGroups: DB_SessionGroup;
33
  sessions: DB_Session;
34
  topics: DB_Topic;
35
  users: DB_User;
36
}
37

38
// Define a local DB
39
export class BrowserDB extends Dexie {
40
  public files: BrowserDBTable<'files'>;
41
  public sessions: BrowserDBTable<'sessions'>;
42
  public messages: BrowserDBTable<'messages'>;
43
  public topics: BrowserDBTable<'topics'>;
44
  public plugins: BrowserDBTable<'plugins'>;
45
  public sessionGroups: BrowserDBTable<'sessionGroups'>;
46
  public users: BrowserDBTable<'users'>;
47

48
  constructor() {
49
    super(LOBE_CHAT_LOCAL_DB_NAME);
50
    this.version(1).stores(dbSchemaV1);
51
    this.version(2).stores(dbSchemaV2);
52
    this.version(3).stores(dbSchemaV3);
53
    this.version(4)
54
      .stores(dbSchemaV4)
55
      .upgrade((trans) => this.upgradeToV4(trans));
56

57
    this.version(5)
58
      .stores(dbSchemaV5)
59
      .upgrade((trans) => this.upgradeToV5(trans));
60

61
    this.version(6)
62
      .stores(dbSchemaV6)
63
      .upgrade((trans) => this.upgradeToV6(trans));
64

65
    this.version(7)
66
      .stores(dbSchemaV7)
67
      .upgrade((trans) => this.upgradeToV7(trans));
68

69
    this.version(8)
70
      .stores(dbSchemaV7)
71
      .upgrade((trans) => this.upgradeToV8(trans));
72

73
    this.version(9)
74
      .stores(dbSchemaV9)
75
      .upgrade((trans) => this.upgradeToV9(trans));
76

77
    this.version(10)
78
      .stores(dbSchemaV9)
79
      .upgrade((trans) => this.upgradeToV10(trans));
80

81
    this.version(11)
82
      .stores(dbSchemaV9)
83
      .upgrade((trans) => this.upgradeToV11(trans));
84

85
    this.files = this.table('files');
86
    this.sessions = this.table('sessions');
87
    this.messages = this.table('messages');
88
    this.topics = this.table('topics');
89
    this.plugins = this.table('plugins');
90
    this.sessionGroups = this.table('sessionGroups');
91
    this.users = this.table('users');
92
  }
93

94
  /**
95
   * 2024.01.22
96
   *
97
   * DB V3 to V4
98
   * from `group = pinned` to `pinned:true`
99
   */
100
  upgradeToV4 = async (trans: Transaction) => {
101
    const sessions = trans.table('sessions');
102
    await sessions.toCollection().modify((session) => {
103
      // translate boolean to number
104
      session.pinned = session.group === 'pinned' ? 1 : 0;
105
      session.group = 'default';
106
    });
107
  };
108

109
  /**
110
   * 2024.01.29
111
   * settings from localStorage to indexedDB
112
   */
113
  upgradeToV5 = async (trans: Transaction) => {
114
    const users = trans.table('users');
115

116
    // if no user, create one
117
    if ((await users.count()) === 0) {
118
      const data = localStorage.getItem('LOBE_SETTINGS');
119

120
      if (data) {
121
        let json;
122

123
        try {
124
          json = JSON.parse(data);
125
        } catch {
126
          /* empty */
127
        }
128

129
        if (!json?.state?.settings) return;
130

131
        const settings = json.state.settings;
132

133
        const user = migrateSettingsToUser(settings);
134
        await users.add(user);
135
      }
136
    }
137
  };
138

139
  /**
140
   * 2024.02.27
141
   * add uuid to user
142
   */
143
  upgradeToV6 = async (trans: Transaction) => {
144
    const users = trans.table('users');
145

146
    await users.toCollection().modify((user: DB_User) => {
147
      if (!user.uuid) user.uuid = uuid();
148
    });
149
  };
150

151
  /**
152
   * 2024.03.14
153
   * add `id` in plugins
154
   */
155
  upgradeToV7 = async (trans: Transaction) => {
156
    const plugins = trans.table('plugins');
157

158
    await plugins.toCollection().modify((plugin: DB_Plugin) => {
159
      plugin.id = plugin.identifier;
160
    });
161
  };
162

163
  upgradeToV8 = async (trans: Transaction) => {
164
    const users = trans.table('users');
165
    await users.toCollection().modify((user: DB_User) => {
166
      if (user.settings) {
167
        user.settings = MigrationLLMSettings.migrateSettings(user.settings as any);
168
      }
169
    });
170
  };
171

172
  /**
173
   * 2024.05.11
174
   *
175
   * message role=function to role=tool
176
   */
177
  upgradeToV9 = async (trans: Transaction) => {
178
    const messages = trans.table('messages');
179
    await messages.toCollection().modify(async (message: DBModel<DB_Message>) => {
180
      if ((message.role as string) === 'function') {
181
        const origin = Object.assign({}, message);
182

183
        const toolCallId = `tool_call_${message.id}`;
184
        const assistantMessageId = `tool_calls_${message.id}`;
185

186
        message.role = 'tool';
187
        message.tool_call_id = toolCallId;
188
        message.parentId = assistantMessageId;
189

190
        await messages.add({
191
          ...origin,
192
          content: '',
193
          createdAt: message.createdAt - 10,
194
          error: undefined,
195
          id: assistantMessageId,
196
          role: 'assistant',
197
          tools: [{ ...message.plugin!, id: toolCallId }],
198
          updatedAt: message.updatedAt - 10,
199
        } as DBModel<DB_Message>);
200
      }
201
    });
202
  };
203

204
  /**
205
   * 2024.05.25
206
   * migrate some agent config to chatConfig
207
   */
208
  upgradeToV10 = async (trans: Transaction) => {
209
    const sessions = trans.table('sessions');
210
    await sessions.toCollection().modify(async (session: DBModel<DB_Session>) => {
211
      if (session.config)
212
        session.config = MigrationAgentChatConfig.migrateChatConfig(session.config as any);
213
    });
214
  };
215

216
  /**
217
   * 2024.05.27
218
   * migrate apiKey in languageModel to keyVaults
219
   */
220
  upgradeToV11 = async (trans: Transaction) => {
221
    const users = trans.table('users');
222

223
    await users.toCollection().modify((user: DB_User) => {
224
      if (user.settings) {
225
        user.settings = MigrationKeyValueSettings.migrateSettings(user.settings as any);
226
      }
227
    });
228
  };
229
}
230

231
export const browserDB = new BrowserDB();
232

233
// ================================================ //
234
// ================================================ //
235
// ================================================ //
236
// ================================================ //
237
// ================================================ //
238

239
// types helper
240
export type BrowserDBSchema = {
241
  [t in keyof LobeDBSchemaMap]: {
242
    model: LobeDBSchemaMap[t];
243
    table: Dexie.Table<DBModel<LobeDBSchemaMap[t]>, string>;
244
  };
245
};
246
type BrowserDBTable<T extends keyof LobeDBSchemaMap> = BrowserDBSchema[T]['table'];
247