lobe-chat

1
import { z } from 'zod';
2

3
import { PluginModel } from '@/database/server/models/plugin';
4
import { authedProcedure, publicProcedure, router } from '@/libs/trpc';
5
import { LobeTool } from '@/types/tool';
6

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

10
  return opts.next({
11
    ctx: { pluginModel: new PluginModel(ctx.userId) },
12
  });
13
});
14

15
export const pluginRouter = router({
16
  createOrInstallPlugin: pluginProcedure
17
    .input(
18
      z.object({
19
        customParams: z.any(),
20
        identifier: z.string(),
21
        manifest: z.any(),
22
        type: z.enum(['plugin', 'customPlugin']),
23
      }),
24
    )
25
    .mutation(async ({ input, ctx }) => {
26
      const result = await ctx.pluginModel.findById(input.identifier);
27

28
      // if not exist, we should create the plugin
29
      if (!result) {
30
        const data = await ctx.pluginModel.create({
31
          customParams: input.customParams,
32
          identifier: input.identifier,
33
          manifest: input.manifest,
34
          type: input.type,
35
        });
36

37
        return data.identifier;
38
      }
39

40
      // or we can just update the plugin manifest
41
      await ctx.pluginModel.update(input.identifier, { manifest: input.manifest });
42
    }),
43

44
  createPlugin: pluginProcedure
45
    .input(
46
      z.object({
47
        customParams: z.any(),
48
        identifier: z.string(),
49
        manifest: z.any(),
50
        type: z.enum(['plugin', 'customPlugin']),
51
      }),
52
    )
53
    .mutation(async ({ input, ctx }) => {
54
      const data = await ctx.pluginModel.create({
55
        customParams: input.customParams,
56
        identifier: input.identifier,
57
        manifest: input.manifest,
58
        type: input.type,
59
      });
60

61
      return data.identifier;
62
    }),
63

64
  getPlugins: publicProcedure.query(async ({ ctx }): Promise<LobeTool[]> => {
65
    if (!ctx.userId) return [];
66

67
    const pluginModel = new PluginModel(ctx.userId);
68

69
    return pluginModel.query();
70
  }),
71

72
  removeAllPlugins: pluginProcedure.mutation(async ({ ctx }) => {
73
    return ctx.pluginModel.deleteAll();
74
  }),
75

76
  removePlugin: pluginProcedure
77
    .input(z.object({ id: z.string() }))
78
    .mutation(async ({ input, ctx }) => {
79
      return ctx.pluginModel.delete(input.id);
80
    }),
81

82
  updatePlugin: pluginProcedure
83
    .input(
84
      z.object({
85
        customParams: z.any().optional(),
86
        id: z.string(),
87
        manifest: z.any().optional(),
88
        settings: z.any().optional(),
89
      }),
90
    )
91
    .mutation(async ({ input, ctx }) => {
92
      return ctx.pluginModel.update(input.id, {
93
        customParams: input.customParams,
94
        manifest: input.manifest,
95
        settings: input.settings,
96
      });
97
    }),
98
});
99

100
export type PluginRouter = typeof pluginRouter;
101