lobe-chat

1
import { z } from 'zod';
2

3
import { KnowledgeBaseModel } from '@/database/server/models/knowledgeBase';
4
import { insertKnowledgeBasesSchema } from '@/database/server/schemas/lobechat';
5
import { authedProcedure, router } from '@/libs/trpc';
6
import { KnowledgeBaseItem } from '@/types/knowledgeBase';
7

8
const knowledgeBaseProcedure = authedProcedure.use(async (opts) => {
9
  const { ctx } = opts;
10

11
  return opts.next({
12
    ctx: {
13
      knowledgeBaseModel: new KnowledgeBaseModel(ctx.userId),
14
    },
15
  });
16
});
17

18
export const knowledgeBaseRouter = router({
19
  addFilesToKnowledgeBase: knowledgeBaseProcedure
20
    .input(z.object({ ids: z.array(z.string()), knowledgeBaseId: z.string() }))
21
    .mutation(async ({ input, ctx }) => {
22
      return ctx.knowledgeBaseModel.addFilesToKnowledgeBase(input.knowledgeBaseId, input.ids);
23
    }),
24

25
  createKnowledgeBase: knowledgeBaseProcedure
26
    .input(
27
      z.object({
28
        avatar: z.string().optional(),
29
        description: z.string().optional(),
30
        name: z.string(),
31
      }),
32
    )
33
    .mutation(async ({ input, ctx }) => {
34
      const data = await ctx.knowledgeBaseModel.create({
35
        avatar: input.avatar,
36
        description: input.description,
37
        name: input.name,
38
      });
39

40
      return data?.id;
41
    }),
42

43
  getKnowledgeBaseById: knowledgeBaseProcedure
44
    .input(z.object({ id: z.string() }))
45
    .query(async ({ ctx, input }): Promise<KnowledgeBaseItem | undefined> => {
46
      return ctx.knowledgeBaseModel.findById(input.id);
47
    }),
48

49
  getKnowledgeBases: knowledgeBaseProcedure.query(async ({ ctx }): Promise<KnowledgeBaseItem[]> => {
50
    return ctx.knowledgeBaseModel.query();
51
  }),
52

53
  removeAllKnowledgeBases: knowledgeBaseProcedure.mutation(async ({ ctx }) => {
54
    return ctx.knowledgeBaseModel.deleteAll();
55
  }),
56

57
  removeFilesFromKnowledgeBase: knowledgeBaseProcedure
58
    .input(z.object({ ids: z.array(z.string()), knowledgeBaseId: z.string() }))
59
    .mutation(async ({ input, ctx }) => {
60
      return ctx.knowledgeBaseModel.removeFilesFromKnowledgeBase(input.knowledgeBaseId, input.ids);
61
    }),
62

63
  removeKnowledgeBase: knowledgeBaseProcedure
64
    .input(z.object({ id: z.string(), removeFiles: z.boolean().optional() }))
65
    .mutation(async ({ input, ctx }) => {
66
      return ctx.knowledgeBaseModel.delete(input.id);
67
    }),
68

69
  updateKnowledgeBase: knowledgeBaseProcedure
70
    .input(
71
      z.object({
72
        id: z.string(),
73
        value: insertKnowledgeBasesSchema.partial(),
74
      }),
75
    )
76
    .mutation(async ({ input, ctx }) => {
77
      return ctx.knowledgeBaseModel.update(input.id, input.value);
78
    }),
79
});
80