lobe-chat

1
import { TRPCError } from '@trpc/server';
2
import { z } from 'zod';
3

4
import { AsyncTaskModel } from '@/database/server/models/asyncTask';
5
import { ChunkModel } from '@/database/server/models/chunk';
6
import { FileModel } from '@/database/server/models/file';
7
import { authedProcedure, router } from '@/libs/trpc';
8
import { S3 } from '@/server/modules/S3';
9
import { getFullFileUrl } from '@/server/utils/files';
10
import { AsyncTaskStatus, AsyncTaskType } from '@/types/asyncTask';
11
import { FileListItem, QueryFileListSchema, UploadFileSchema } from '@/types/files';
12

13
const fileProcedure = authedProcedure.use(async (opts) => {
14
  const { ctx } = opts;
15

16
  return opts.next({
17
    ctx: {
18
      asyncTaskModel: new AsyncTaskModel(ctx.userId),
19
      chunkModel: new ChunkModel(ctx.userId),
20
      fileModel: new FileModel(ctx.userId),
21
    },
22
  });
23
});
24

25
export const fileRouter = router({
26
  checkFileHash: fileProcedure
27
    .input(z.object({ hash: z.string() }))
28
    .mutation(async ({ ctx, input }) => {
29
      return ctx.fileModel.checkHash(input.hash);
30
    }),
31

32
  createFile: fileProcedure
33
    .input(
34
      UploadFileSchema.omit({ data: true, saveMode: true, url: true }).extend({ url: z.string() }),
35
    )
36
    .mutation(async ({ ctx, input }) => {
37
      const { isExist } = await ctx.fileModel.checkHash(input.hash!);
38

39
      // if the file is not exist in global file, create a new one
40
      if (!isExist) {
41
        await ctx.fileModel.createGlobalFile({
42
          fileType: input.fileType,
43
          hashId: input.hash!,
44
          metadata: input.metadata,
45
          size: input.size,
46
          url: input.url,
47
        });
48
      }
49

50
      const { id } = await ctx.fileModel.create({
51
        fileHash: input.hash,
52
        fileType: input.fileType,
53
        knowledgeBaseId: input.knowledgeBaseId,
54
        metadata: input.metadata,
55
        name: input.name,
56
        size: input.size,
57
        url: input.url,
58
      });
59

60
      return { id, url: getFullFileUrl(input.url) };
61
    }),
62
  findById: fileProcedure
63
    .input(
64
      z.object({
65
        id: z.string(),
66
      }),
67
    )
68
    .query(async ({ ctx, input }) => {
69
      const item = await ctx.fileModel.findById(input.id);
70
      if (!item) throw new TRPCError({ code: 'BAD_REQUEST', message: 'File not found' });
71

72
      return { ...item, url: getFullFileUrl(item?.url) };
73
    }),
74

75
  getFileItemById: fileProcedure
76
    .input(
77
      z.object({
78
        id: z.string(),
79
      }),
80
    )
81
    .query(async ({ ctx, input }): Promise<FileListItem | undefined> => {
82
      const item = await ctx.fileModel.findById(input.id);
83

84
      if (!item) throw new TRPCError({ code: 'BAD_REQUEST', message: 'File not found' });
85

86
      let embeddingTask = null;
87
      if (item.embeddingTaskId) {
88
        embeddingTask = await ctx.asyncTaskModel.findById(item.embeddingTaskId);
89
      }
90
      let chunkingTask = null;
91
      if (item.chunkTaskId) {
92
        chunkingTask = await ctx.asyncTaskModel.findById(item.chunkTaskId);
93
      }
94

95
      const chunkCount = await ctx.chunkModel.countByFileId(input.id);
96

97
      return {
98
        ...item,
99
        chunkCount,
100
        chunkingError: chunkingTask?.error,
101
        chunkingStatus: chunkingTask?.status as AsyncTaskStatus,
102
        embeddingError: embeddingTask?.error,
103
        embeddingStatus: embeddingTask?.status as AsyncTaskStatus,
104
        finishEmbedding: embeddingTask?.status === AsyncTaskStatus.Success,
105
        url: getFullFileUrl(item.url!),
106
      };
107
    }),
108

109
  getFiles: fileProcedure.input(QueryFileListSchema).query(async ({ ctx, input }) => {
110
    const fileList = await ctx.fileModel.query(input);
111

112
    const fileIds = fileList.map((item) => item.id);
113
    const chunks = await ctx.chunkModel.countByFileIds(fileIds);
114

115
    const chunkTaskIds = fileList.map((result) => result.chunkTaskId).filter(Boolean) as string[];
116

117
    const chunkTasks = await ctx.asyncTaskModel.findByIds(chunkTaskIds, AsyncTaskType.Chunking);
118

119
    const embeddingTaskIds = fileList
120
      .map((result) => result.embeddingTaskId)
121
      .filter(Boolean) as string[];
122
    const embeddingTasks = await ctx.asyncTaskModel.findByIds(
123
      embeddingTaskIds,
124
      AsyncTaskType.Embedding,
125
    );
126

127
    return fileList.map(({ chunkTaskId, embeddingTaskId, ...item }): FileListItem => {
128
      const chunkTask = chunkTaskId ? chunkTasks.find((task) => task.id === chunkTaskId) : null;
129
      const embeddingTask = embeddingTaskId
130
        ? embeddingTasks.find((task) => task.id === embeddingTaskId)
131
        : null;
132

133
      return {
134
        ...item,
135
        chunkCount: chunks.find((chunk) => chunk.id === item.id)?.count ?? null,
136
        chunkingError: chunkTask?.error ?? null,
137
        chunkingStatus: chunkTask?.status as AsyncTaskStatus,
138
        embeddingError: embeddingTask?.error ?? null,
139
        embeddingStatus: embeddingTask?.status as AsyncTaskStatus,
140
        finishEmbedding: embeddingTask?.status === AsyncTaskStatus.Success,
141
        url: getFullFileUrl(item.url!),
142
      };
143
    });
144
  }),
145

146
  removeAllFiles: fileProcedure.mutation(async ({ ctx }) => {
147
    return ctx.fileModel.clear();
148
  }),
149

150
  removeFile: fileProcedure.input(z.object({ id: z.string() })).mutation(async ({ input, ctx }) => {
151
    const file = await ctx.fileModel.delete(input.id);
152

153
    // delete the orphan chunks
154
    await ctx.chunkModel.deleteOrphanChunks();
155
    if (!file) return;
156

157
    // delele the file from remove from S3 if it is not used by other files
158
    const s3Client = new S3();
159
    await s3Client.deleteFile(file.url!);
160
  }),
161

162
  removeFileAsyncTask: fileProcedure
163
    .input(
164
      z.object({
165
        id: z.string(),
166
        type: z.enum(['embedding', 'chunk']),
167
      }),
168
    )
169
    .mutation(async ({ ctx, input }) => {
170
      const file = await ctx.fileModel.findById(input.id);
171

172
      if (!file) return;
173

174
      const taskId = input.type === 'embedding' ? file.embeddingTaskId : file.chunkTaskId;
175

176
      if (!taskId) return;
177

178
      await ctx.asyncTaskModel.delete(taskId);
179
    }),
180

181
  removeFiles: fileProcedure
182
    .input(z.object({ ids: z.array(z.string()) }))
183
    .mutation(async ({ input, ctx }) => {
184
      const needToRemoveFileList = await ctx.fileModel.deleteMany(input.ids);
185

186
      // delete the orphan chunks
187
      await ctx.chunkModel.deleteOrphanChunks();
188

189
      if (!needToRemoveFileList || needToRemoveFileList.length === 0) return;
190

191
      // remove from S3
192
      const s3Client = new S3();
193

194
      await s3Client.deleteFiles(needToRemoveFileList.map((file) => file.url!));
195
    }),
196
});
197

198
export type FileRouter = typeof fileRouter;
199