lobe-chat

1
import { z } from 'zod';
2

3
import { INBOX_SESSION_ID } from '@/const/session';
4
import { DEFAULT_AGENT_CONFIG } from '@/const/settings';
5
import { AgentModel } from '@/database/server/models/agent';
6
import { FileModel } from '@/database/server/models/file';
7
import { KnowledgeBaseModel } from '@/database/server/models/knowledgeBase';
8
import { SessionModel } from '@/database/server/models/session';
9
import { UserModel } from '@/database/server/models/user';
10
import { pino } from '@/libs/logger';
11
import { authedProcedure, router } from '@/libs/trpc';
12
import { KnowledgeItem, KnowledgeType } from '@/types/knowledgeBase';
13

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

17
  return opts.next({
18
    ctx: {
19
      agentModel: new AgentModel(ctx.userId),
20
      fileModel: new FileModel(ctx.userId),
21
      knowledgeBaseModel: new KnowledgeBaseModel(ctx.userId),
22
      sessionModel: new SessionModel(ctx.userId),
23
    },
24
  });
25
});
26

27
export const agentRouter = router({
28
  createAgentFiles: agentProcedure
29
    .input(
30
      z.object({
31
        agentId: z.string(),
32
        enabled: z.boolean().optional(),
33
        fileIds: z.array(z.string()),
34
      }),
35
    )
36
    .mutation(async ({ input, ctx }) => {
37
      return ctx.agentModel.createAgentFiles(input.agentId, input.fileIds, input.enabled);
38
    }),
39

40
  createAgentKnowledgeBase: agentProcedure
41
    .input(
42
      z.object({
43
        agentId: z.string(),
44
        enabled: z.boolean().optional(),
45
        knowledgeBaseId: z.string(),
46
      }),
47
    )
48
    .mutation(async ({ input, ctx }) => {
49
      return ctx.agentModel.createAgentKnowledgeBase(
50
        input.agentId,
51
        input.knowledgeBaseId,
52
        input.enabled,
53
      );
54
    }),
55

56
  deleteAgentFile: agentProcedure
57
    .input(
58
      z.object({
59
        agentId: z.string(),
60
        fileId: z.string(),
61
      }),
62
    )
63
    .mutation(async ({ input, ctx }) => {
64
      return ctx.agentModel.deleteAgentFile(input.agentId, input.fileId);
65
    }),
66

67
  deleteAgentKnowledgeBase: agentProcedure
68
    .input(
69
      z.object({
70
        agentId: z.string(),
71
        knowledgeBaseId: z.string(),
72
      }),
73
    )
74
    .mutation(async ({ input, ctx }) => {
75
      return ctx.agentModel.deleteAgentKnowledgeBase(input.agentId, input.knowledgeBaseId);
76
    }),
77

78
  getAgentConfig: agentProcedure
79
    .input(
80
      z.object({
81
        sessionId: z.string(),
82
      }),
83
    )
84
    .query(async ({ input, ctx }) => {
85
      if (input.sessionId === INBOX_SESSION_ID) {
86
        const item = await ctx.sessionModel.findByIdOrSlug(INBOX_SESSION_ID);
87
        // if there is no session for user, create one
88
        if (!item) {
89
          // if there is no user, return default config
90
          const user = await UserModel.findById(ctx.userId);
91
          if (!user) return DEFAULT_AGENT_CONFIG;
92

93
          const res = await ctx.sessionModel.createInbox();
94
          pino.info('create inbox session', res);
95
        }
96
      }
97

98
      const session = await ctx.sessionModel.findByIdOrSlug(input.sessionId);
99

100
      if (!session) throw new Error('Session not found');
101
      const sessionId = session.id;
102

103
      return await ctx.agentModel.findBySessionId(sessionId);
104
    }),
105

106
  getKnowledgeBasesAndFiles: agentProcedure
107
    .input(
108
      z.object({
109
        agentId: z.string(),
110
      }),
111
    )
112
    .query(async ({ ctx, input }): Promise<KnowledgeItem[]> => {
113
      const knowledgeBases = await ctx.knowledgeBaseModel.query();
114

115
      const files = await ctx.fileModel.query({
116
        showFilesInKnowledgeBase: false,
117
      });
118

119
      const knowledge = await ctx.agentModel.getAgentAssignedKnowledge(input.agentId);
120

121
      return [
122
        ...files.map((file) => ({
123
          enabled: knowledge.files.some((item) => item.id === file.id),
124
          fileType: file.fileType,
125
          id: file.id,
126
          name: file.name,
127
          type: KnowledgeType.File,
128
        })),
129
        ...knowledgeBases.map((knowledgeBase) => ({
130
          avatar: knowledgeBase.avatar,
131
          description: knowledgeBase.description,
132
          enabled: knowledge.knowledgeBases.some((item) => item.id === knowledgeBase.id),
133
          id: knowledgeBase.id,
134
          name: knowledgeBase.name,
135
          type: KnowledgeType.KnowledgeBase,
136
        })),
137
      ];
138
    }),
139

140
  toggleFile: agentProcedure
141
    .input(
142
      z.object({
143
        agentId: z.string(),
144
        enabled: z.boolean().optional(),
145
        fileId: z.string(),
146
      }),
147
    )
148
    .mutation(async ({ input, ctx }) => {
149
      return ctx.agentModel.toggleFile(input.agentId, input.fileId, input.enabled);
150
    }),
151

152
  toggleKnowledgeBase: agentProcedure
153
    .input(
154
      z.object({
155
        agentId: z.string(),
156
        enabled: z.boolean().optional(),
157
        knowledgeBaseId: z.string(),
158
      }),
159
    )
160
    .mutation(async ({ input, ctx }) => {
161
      return ctx.agentModel.toggleKnowledgeBase(
162
        input.agentId,
163
        input.knowledgeBaseId,
164
        input.enabled,
165
      );
166
    }),
167
});
168