lobe-chat

1
import { z } from 'zod';
2

3
import { SessionModel } from '@/database/server/models/session';
4
import { SessionGroupModel } from '@/database/server/models/sessionGroup';
5
import { insertAgentSchema, insertSessionSchema } from '@/database/server/schemas/lobechat';
6
import { authedProcedure, publicProcedure, router } from '@/libs/trpc';
7
import { AgentChatConfigSchema } from '@/types/agent';
8
import { LobeMetaDataSchema } from '@/types/meta';
9
import { BatchTaskResult } from '@/types/service';
10
import { ChatSessionList } from '@/types/session';
11
import { merge } from '@/utils/merge';
12

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

16
  return opts.next({
17
    ctx: {
18
      sessionGroupModel: new SessionGroupModel(ctx.userId),
19
      sessionModel: new SessionModel(ctx.userId),
20
    },
21
  });
22
});
23

24
export const sessionRouter = router({
25
  batchCreateSessions: sessionProcedure
26
    .input(
27
      z.array(
28
        z
29
          .object({
30
            config: z.object({}).passthrough(),
31
            group: z.string().optional(),
32
            id: z.string(),
33
            meta: LobeMetaDataSchema,
34
            pinned: z.boolean().optional(),
35
            type: z.string(),
36
          })
37
          .partial(),
38
      ),
39
    )
40
    .mutation(async ({ input, ctx }): Promise<BatchTaskResult> => {
41
      const data = await ctx.sessionModel.batchCreate(
42
        input.map((item) => ({
43
          ...item,
44
          ...item.meta,
45
        })) as any,
46
      );
47

48
      return { added: data.rowCount as number, ids: [], skips: [], success: true };
49
    }),
50

51
  cloneSession: sessionProcedure
52
    .input(z.object({ id: z.string(), newTitle: z.string() }))
53
    .mutation(async ({ input, ctx }) => {
54
      const data = await ctx.sessionModel.duplicate(input.id, input.newTitle);
55

56
      return data?.id;
57
    }),
58

59
  countSessions: sessionProcedure.query(async ({ ctx }) => {
60
    return ctx.sessionModel.count();
61
  }),
62

63
  createSession: sessionProcedure
64
    .input(
65
      z.object({
66
        config: insertAgentSchema
67
          .omit({ chatConfig: true, plugins: true, tags: true, tts: true })
68
          .passthrough()
69
          .partial(),
70
        session: insertSessionSchema.omit({ createdAt: true, updatedAt: true }).partial(),
71
        type: z.enum(['agent', 'group']),
72
      }),
73
    )
74
    .mutation(async ({ input, ctx }) => {
75
      const data = await ctx.sessionModel.create(input);
76

77
      return data.id;
78
    }),
79

80
  getGroupedSessions: publicProcedure.query(async ({ ctx }): Promise<ChatSessionList> => {
81
    if (!ctx.userId)
82
      return {
83
        sessionGroups: [],
84
        sessions: [],
85
      };
86

87
    const sessionModel = new SessionModel(ctx.userId);
88

89
    return sessionModel.queryWithGroups();
90
  }),
91

92
  getSessions: sessionProcedure
93
    .input(
94
      z.object({
95
        current: z.number().optional(),
96
        pageSize: z.number().optional(),
97
      }),
98
    )
99
    .query(async ({ input, ctx }) => {
100
      const { current, pageSize } = input;
101

102
      return ctx.sessionModel.query({ current, pageSize });
103
    }),
104

105
  removeAllSessions: sessionProcedure.mutation(async ({ ctx }) => {
106
    return ctx.sessionModel.deleteAll();
107
  }),
108

109
  removeSession: sessionProcedure
110
    .input(z.object({ id: z.string() }))
111
    .mutation(async ({ input, ctx }) => {
112
      return ctx.sessionModel.delete(input.id);
113
    }),
114

115
  searchSessions: sessionProcedure
116
    .input(z.object({ keywords: z.string() }))
117
    .query(async ({ input, ctx }) => {
118
      return ctx.sessionModel.queryByKeyword(input.keywords);
119
    }),
120

121
  updateSession: sessionProcedure
122
    .input(
123
      z.object({
124
        id: z.string(),
125
        value: insertSessionSchema.partial(),
126
      }),
127
    )
128
    .mutation(async ({ input, ctx }) => {
129
      return ctx.sessionModel.update(input.id, input.value);
130
    }),
131
  updateSessionChatConfig: sessionProcedure
132
    .input(
133
      z.object({
134
        id: z.string(),
135
        value: AgentChatConfigSchema.partial(),
136
      }),
137
    )
138
    .mutation(async ({ input, ctx }) => {
139
      const session = await ctx.sessionModel.findByIdOrSlug(input.id);
140

141
      if (!session) return;
142

143
      return ctx.sessionModel.updateConfig(session.agent.id, {
144
        chatConfig: merge(session.agent.chatConfig, input.value),
145
      });
146
    }),
147
  updateSessionConfig: sessionProcedure
148
    .input(
149
      z.object({
150
        id: z.string(),
151
        value: z.object({}).passthrough().partial(),
152
      }),
153
    )
154
    .mutation(async ({ input, ctx }) => {
155
      const session = await ctx.sessionModel.findByIdOrSlug(input.id);
156

157
      if (!session || !input.value) return;
158

159
      if (!session.agent) {
160
        throw new Error(
161
          'this session is not assign with agent, please contact with admin to fix this issue.',
162
        );
163
      }
164

165
      return ctx.sessionModel.updateConfig(session.agent.id, input.value);
166
    }),
167
});
168

169
export type SessionRouter = typeof sessionRouter;
170