lobe-chat

1
import { DefaultErrorShape } from '@trpc/server/unstable-core-do-not-import';
2

3
import { edgeClient, lambdaClient } from '@/libs/trpc/client';
4
import { useUserStore } from '@/store/user';
5
import { ImportStage, ImporterEntryData, OnImportCallbacks } from '@/types/importer';
6
import { UserSettings } from '@/types/user/settings';
7
import { uuid } from '@/utils/uuid';
8

9
export class ServerService {
10
  importSettings = async (settings: UserSettings) => {
11
    await useUserStore.getState().importAppSettings(settings);
12
  };
13

14
  importData = async (data: ImporterEntryData, callbacks?: OnImportCallbacks): Promise<void> => {
15
    const handleError = (e: unknown) => {
16
      callbacks?.onStageChange?.(ImportStage.Error);
17
      const error = e as DefaultErrorShape;
18

19
      callbacks?.onError?.({
20
        code: error.data.code,
21
        httpStatus: error.data.httpStatus,
22
        message: error.message,
23
        path: error.data.path,
24
      });
25
    };
26

27
    const totalLength =
28
      (data.messages?.length || 0) +
29
      (data.sessionGroups?.length || 0) +
30
      (data.sessions?.length || 0) +
31
      (data.topics?.length || 0);
32

33
    if (totalLength < 500) {
34
      callbacks?.onStageChange?.(ImportStage.Importing);
35
      const time = Date.now();
36
      try {
37
        const result = await lambdaClient.importer.importByPost.mutate({ data });
38
        const duration = Date.now() - time;
39

40
        callbacks?.onStageChange?.(ImportStage.Success);
41
        callbacks?.onSuccess?.(result, duration);
42
      } catch (e) {
43
        handleError(e);
44
      }
45

46
      return;
47
    }
48

49
    // if the data is too large, upload it to S3 and upload by file
50
    const filename = `${uuid()}.json`;
51

52
    const pathname = `import_config/${filename}`;
53

54
    const url = await edgeClient.upload.createS3PreSignedUrl.mutate({ pathname });
55

56
    try {
57
      callbacks?.onStageChange?.(ImportStage.Uploading);
58
      await this.uploadWithProgress(url, data, callbacks?.onFileUploading);
59
    } catch {
60
      throw new Error('Upload Error');
61
    }
62

63
    callbacks?.onStageChange?.(ImportStage.Importing);
64
    const time = Date.now();
65
    try {
66
      const result = await lambdaClient.importer.importByFile.mutate({ pathname });
67
      const duration = Date.now() - time;
68
      callbacks?.onStageChange?.(ImportStage.Success);
69
      callbacks?.onSuccess?.(result, duration);
70
    } catch (e) {
71
      handleError(e);
72
    }
73
  };
74

75
  private uploadWithProgress = async (
76
    url: string,
77
    data: object,
78
    onProgress: OnImportCallbacks['onFileUploading'],
79
  ) => {
80
    const xhr = new XMLHttpRequest();
81

82
    let startTime = Date.now();
83
    xhr.upload.addEventListener('progress', (event) => {
84
      if (event.lengthComputable) {
85
        const progress = Number(((event.loaded / event.total) * 100).toFixed(1));
86

87
        const speedInByte = event.loaded / ((Date.now() - startTime) / 1000);
88

89
        onProgress?.({
90
          // if the progress is 100, it means the file is uploaded
91
          // but the server is still processing it
92
          // so make it as 99.5 and let users think it's still uploading
93
          progress: progress === 100 ? 99.5 : progress,
94
          restTime: (event.total - event.loaded) / speedInByte,
95
          speed: speedInByte / 1024,
96
        });
97
      }
98
    });
99

100
    xhr.open('PUT', url);
101
    xhr.setRequestHeader('Content-Type', 'application/json');
102

103
    return new Promise((resolve, reject) => {
104
      xhr.addEventListener('load', () => {
105
        if (xhr.status >= 200 && xhr.status < 300) {
106
          resolve(xhr.response);
107
        } else {
108
          reject(xhr.statusText);
109
        }
110
      });
111
      xhr.addEventListener('error', () => reject(xhr.statusText));
112
      xhr.send(JSON.stringify(data));
113
    });
114
  };
115
}
116