lobe-chat

1
import { z } from 'zod';
2

3
import { MessageModel } from '@/database/server/models/message';
4
import { updateMessagePluginSchema } from '@/database/server/schemas/lobechat';
5
import { authedProcedure, publicProcedure, router } from '@/libs/trpc';
6
import { ChatMessage } from '@/types/message';
7
import { BatchTaskResult } from '@/types/service';
8

9
type ChatMessageList = ChatMessage[];
10

11
const messageProcedure = authedProcedure.use(async (opts) => {
12
  const { ctx } = opts;
13

14
  return opts.next({
15
    ctx: { messageModel: new MessageModel(ctx.userId) },
16
  });
17
});
18

19
export const messageRouter = router({
20
  batchCreateMessages: messageProcedure
21
    .input(z.array(z.any()))
22
    .mutation(async ({ input, ctx }): Promise<BatchTaskResult> => {
23
      const data = await ctx.messageModel.batchCreate(input);
24

25
      return { added: data.rowCount as number, ids: [], skips: [], success: true };
26
    }),
27

28
  count: messageProcedure.query(async ({ ctx }) => {
29
    return ctx.messageModel.count();
30
  }),
31
  countToday: messageProcedure.query(async ({ ctx }) => {
32
    return ctx.messageModel.countToday();
33
  }),
34

35
  createMessage: messageProcedure
36
    .input(z.object({}).passthrough().partial())
37
    .mutation(async ({ input, ctx }) => {
38
      const data = await ctx.messageModel.create(input as any);
39

40
      return data.id;
41
    }),
42

43
  getAllMessages: messageProcedure.query(async ({ ctx }): Promise<ChatMessageList> => {
44
    return ctx.messageModel.queryAll();
45
  }),
46

47
  getAllMessagesInSession: messageProcedure
48
    .input(
49
      z.object({
50
        sessionId: z.string().nullable().optional(),
51
      }),
52
    )
53
    .query(async ({ ctx, input }): Promise<ChatMessageList> => {
54
      return ctx.messageModel.queryBySessionId(input.sessionId);
55
    }),
56

57
  getMessages: publicProcedure
58
    .input(
59
      z.object({
60
        current: z.number().optional(),
61
        pageSize: z.number().optional(),
62
        sessionId: z.string().nullable().optional(),
63
        topicId: z.string().nullable().optional(),
64
      }),
65
    )
66
    .query(async ({ input, ctx }) => {
67
      if (!ctx.userId) return [];
68

69
      const messageModel = new MessageModel(ctx.userId);
70

71
      return messageModel.query(input);
72
    }),
73

74
  removeAllMessages: messageProcedure.mutation(async ({ ctx }) => {
75
    return ctx.messageModel.deleteAllMessages();
76
  }),
77

78
  removeMessage: messageProcedure
79
    .input(z.object({ id: z.string() }))
80
    .mutation(async ({ input, ctx }) => {
81
      return ctx.messageModel.deleteMessage(input.id);
82
    }),
83

84
  removeMessageQuery: messageProcedure
85
    .input(z.object({ id: z.string() }))
86
    .mutation(async ({ input, ctx }) => {
87
      return ctx.messageModel.deleteMessageQuery(input.id);
88
    }),
89

90
  removeMessages: messageProcedure
91
    .input(z.object({ ids: z.array(z.string()) }))
92
    .mutation(async ({ input, ctx }) => {
93
      return ctx.messageModel.deleteMessages(input.ids);
94
    }),
95

96
  removeMessagesByAssistant: messageProcedure
97
    .input(
98
      z.object({
99
        sessionId: z.string().nullable().optional(),
100
        topicId: z.string().nullable().optional(),
101
      }),
102
    )
103
    .mutation(async ({ input, ctx }) => {
104
      return ctx.messageModel.deleteMessagesBySession(input.sessionId, input.topicId);
105
    }),
106

107
  searchMessages: messageProcedure
108
    .input(z.object({ keywords: z.string() }))
109
    .query(async ({ input, ctx }) => {
110
      return ctx.messageModel.queryByKeyword(input.keywords);
111
    }),
112

113
  update: messageProcedure
114
    .input(
115
      z.object({
116
        id: z.string(),
117
        value: z.object({}).passthrough().partial(),
118
      }),
119
    )
120
    .mutation(async ({ input, ctx }) => {
121
      return ctx.messageModel.update(input.id, input.value);
122
    }),
123

124
  updateMessagePlugin: messageProcedure
125
    .input(
126
      z.object({
127
        id: z.string(),
128
        value: updateMessagePluginSchema.partial(),
129
      }),
130
    )
131
    .mutation(async ({ input, ctx }) => {
132
      return ctx.messageModel.updateMessagePlugin(input.id, input.value);
133
    }),
134

135
  updatePluginState: messageProcedure
136
    .input(
137
      z.object({
138
        id: z.string(),
139
        value: z.object({}).passthrough(),
140
      }),
141
    )
142
    .mutation(async ({ input, ctx }) => {
143
      return ctx.messageModel.updatePluginState(input.id, input.value);
144
    }),
145

146
  updateTTS: messageProcedure
147
    .input(
148
      z.object({
149
        id: z.string(),
150
        value: z
151
          .object({
152
            contentMd5: z.string().optional(),
153
            file: z.string().optional(),
154
            voice: z.string().optional(),
155
          })
156
          .or(z.literal(false)),
157
      }),
158
    )
159
    .mutation(async ({ input, ctx }) => {
160
      if (input.value === false) {
161
        return ctx.messageModel.deleteMessageTTS(input.id);
162
      }
163

164
      return ctx.messageModel.updateTTS(input.id, input.value);
165
    }),
166

167
  updateTranslate: messageProcedure
168
    .input(
169
      z.object({
170
        id: z.string(),
171
        value: z
172
          .object({
173
            content: z.string().optional(),
174
            from: z.string().optional(),
175
            to: z.string(),
176
          })
177
          .or(z.literal(false)),
178
      }),
179
    )
180
    .mutation(async ({ input, ctx }) => {
181
      if (input.value === false) {
182
        return ctx.messageModel.deleteMessageTranslate(input.id);
183
      }
184

185
      return ctx.messageModel.updateTranslate(input.id, input.value);
186
    }),
187
});
188

189
export type MessageRouter = typeof messageRouter;
190