lobe-chat

1
import { z } from 'zod';
2

3
import { TopicModel } from '@/database/server/models/topic';
4
import { authedProcedure, publicProcedure, router } from '@/libs/trpc';
5
import { BatchTaskResult } from '@/types/service';
6

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

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

15
export const topicRouter = router({
16
  batchCreateTopics: topicProcedure
17
    .input(
18
      z.array(
19
        z.object({
20
          favorite: z.boolean().optional(),
21
          id: z.string().optional(),
22
          messages: z.array(z.string()).optional(),
23
          sessionId: z.string().optional(),
24
          title: z.string(),
25
        }),
26
      ),
27
    )
28
    .mutation(async ({ input, ctx }): Promise<BatchTaskResult> => {
29
      const data = await ctx.topicModel.batchCreate(
30
        input.map((item) => ({
31
          ...item,
32
        })) as any,
33
      );
34

35
      return { added: data.length, ids: [], skips: [], success: true };
36
    }),
37

38
  batchDelete: topicProcedure
39
    .input(z.object({ ids: z.array(z.string()) }))
40
    .mutation(async ({ input, ctx }) => {
41
      return ctx.topicModel.batchDelete(input.ids);
42
    }),
43

44
  batchDeleteBySessionId: topicProcedure
45
    .input(z.object({ id: z.string().nullable().optional() }))
46
    .mutation(async ({ input, ctx }) => {
47
      return ctx.topicModel.batchDeleteBySessionId(input.id);
48
    }),
49

50
  cloneTopic: topicProcedure
51
    .input(z.object({ id: z.string(), newTitle: z.string().optional() }))
52
    .mutation(async ({ input, ctx }) => {
53
      const data = await ctx.topicModel.duplicate(input.id, input.newTitle);
54

55
      return data.topic.id;
56
    }),
57

58
  countTopics: topicProcedure.query(async ({ ctx }) => {
59
    return ctx.topicModel.count();
60
  }),
61

62
  createTopic: topicProcedure
63
    .input(
64
      z.object({
65
        favorite: z.boolean().optional(),
66
        messages: z.array(z.string()).optional(),
67
        sessionId: z.string().nullable().optional(),
68
        title: z.string(),
69
      }),
70
    )
71
    .mutation(async ({ input, ctx }) => {
72
      const data = await ctx.topicModel.create(input);
73

74
      return data.id;
75
    }),
76

77
  getAllTopics: topicProcedure.query(async ({ ctx }) => {
78
    return ctx.topicModel.queryAll();
79
  }),
80

81
  getTopics: publicProcedure
82
    .input(
83
      z.object({
84
        current: z.number().optional(),
85
        pageSize: z.number().optional(),
86
        sessionId: z.string().nullable().optional(),
87
      }),
88
    )
89
    .query(async ({ input, ctx }) => {
90
      if (!ctx.userId) return [];
91

92
      const topicModel = new TopicModel(ctx.userId);
93

94
      return topicModel.query(input);
95
    }),
96

97
  hasTopics: topicProcedure.query(async ({ ctx }) => {
98
    return (await ctx.topicModel.count()) === 0;
99
  }),
100

101
  removeAllTopics: topicProcedure.mutation(async ({ ctx }) => {
102
    return ctx.topicModel.deleteAll();
103
  }),
104

105
  removeTopic: topicProcedure
106
    .input(z.object({ id: z.string() }))
107
    .mutation(async ({ input, ctx }) => {
108
      return ctx.topicModel.delete(input.id);
109
    }),
110

111
  searchTopics: topicProcedure
112
    .input(z.object({ keywords: z.string(), sessionId: z.string().nullable().optional() }))
113
    .query(async ({ input, ctx }) => {
114
      return ctx.topicModel.queryByKeyword(input.keywords, input.sessionId);
115
    }),
116

117
  updateTopic: topicProcedure
118
    .input(
119
      z.object({
120
        id: z.string(),
121
        value: z.object({
122
          favorite: z.boolean().optional(),
123
          messages: z.array(z.string()).optional(),
124
          sessionId: z.string().optional(),
125
          title: z.string().optional(),
126
        }),
127
      }),
128
    )
129
    .mutation(async ({ input, ctx }) => {
130
      return ctx.topicModel.update(input.id, input.value);
131
    }),
132
});
133

134
export type TopicRouter = typeof topicRouter;
135