Skip to content

Commit 6358383

Browse files
danny-aviladanny-avila
authored
feat(db & e2e): Enhance DB Schemas/Controllers and Improve E2E Tests (#966)
* feat: add global teardown to remove test data and add registration/log-out to auth flow * refactor(models/Conversation): index user field and add JSDoc to deleteConvos * refactor: add user index to message schema and ensure user is saved to each Message * refactor: add user to each saveMessage call * fix: handle case where title is null in zod schema * feat(e2e): ensure messages are deleted on cleanUp * fix: set last convo for all endpoints on conversation update * fix: enable registration for CI env
1 parent fd70e21 commit 6358383

28 files changed

+229
-72
lines changed

‎api/app/clients/BaseClient.js

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -515,7 +515,7 @@ class BaseClient {
515515
}
516516

517517
async saveMessageToDatabase(message, endpointOptions, user = null) {
518-
await saveMessage({ ...message, unfinished: false, cancelled: false });
518+
await saveMessage({ ...message, user, unfinished: false, cancelled: false });
519519
await saveConvo(user, {
520520
conversationId: message.conversationId,
521521
endpoint: this.options.endpoint,

‎api/models/Conversation.js

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,23 @@ module.exports = {
108108
return { message: 'Error getting conversation title' };
109109
}
110110
},
111+
/**
112+
* Asynchronously deletes conversations and associated messages for a given user and filter.
113+
*
114+
* @async
115+
* @function
116+
* @param {string|ObjectId} user - The user's ID.
117+
* @param {Object} filter - Additional filter criteria for the conversations to be deleted.
118+
* @returns {Promise<{ n: number, ok: number, deletedCount: number, messages: { n: number, ok: number, deletedCount: number } }>}
119+
* An object containing the count of deleted conversations and associated messages.
120+
* @throws {Error} Throws an error if there's an issue with the database operations.
121+
*
122+
* @example
123+
* const user = 'someUserId';
124+
* const filter = { someField: 'someValue' };
125+
* const result = await deleteConvos(user, filter);
126+
* console.log(result); // { n: 5, ok: 1, deletedCount: 5, messages: { n: 10, ok: 1, deletedCount: 10 } }
127+
*/
111128
deleteConvos: async (user, filter) => {
112129
let toRemove = await Conversation.find({ ...filter, user }).select('conversationId');
113130
const ids = toRemove.map((instance) => instance.conversationId);

‎api/models/Message.js

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ module.exports = {
77
Message,
88

99
async saveMessage({
10+
user,
1011
messageId,
1112
newMessageId,
1213
conversationId,
@@ -33,6 +34,7 @@ module.exports = {
3334
await Message.findOneAndUpdate(
3435
{ messageId },
3536
{
37+
user,
3638
messageId: newMessageId || messageId,
3739
conversationId,
3840
parentMessageId,

‎api/models/schema/convoSchema.js

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ const convoSchema = mongoose.Schema(
1717
},
1818
user: {
1919
type: String,
20+
index: true,
2021
default: null,
2122
},
2223
messages: [{ type: mongoose.Schema.Types.ObjectId, ref: 'Message' }],

‎api/models/schema/messageSchema.js

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,11 @@ const messageSchema = mongoose.Schema(
1414
required: true,
1515
meiliIndex: true,
1616
},
17+
user: {
18+
type: String,
19+
index: true,
20+
default: null,
21+
},
1722
model: {
1823
type: String,
1924
},

‎api/server/middleware/abortMiddleware.js

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ const createAbortController = (req, res, getAbortData) => {
5454
isCreatedByUser: false,
5555
};
5656

57-
saveMessage(responseMessage);
57+
saveMessage({ ...responseMessage, user: req.user.id });
5858

5959
return {
6060
title: await getConvoTitle(req.user.id, conversationId),
@@ -80,6 +80,7 @@ const handleAbortError = async (res, req, error, data) => {
8080
parentMessageId,
8181
text: error.message,
8282
shouldSaveMessage: true,
83+
user: req.user.id,
8384
};
8485
const callback = async () => {
8586
if (abortControllers.has(conversationId)) {

‎api/server/middleware/denyRequest.js

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ const denyRequest = async (req, res, errorMessage) => {
4242
_convoId && parentMessageId && parentMessageId !== '00000000-0000-0000-0000-000000000000';
4343

4444
if (shouldSaveMessage) {
45-
await saveMessage(userMessage);
45+
await saveMessage({ ...userMessage, user: req.user.id });
4646
}
4747

4848
return await sendError(res, {
@@ -52,6 +52,7 @@ const denyRequest = async (req, res, errorMessage) => {
5252
parentMessageId: userMessage.messageId,
5353
text: responseText,
5454
shouldSaveMessage,
55+
user: req.user.id,
5556
});
5657
};
5758

‎api/server/routes/ask/anthropic.js

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
3030
let responseMessageId;
3131
let lastSavedTimestamp = 0;
3232
let saveDelay = 100;
33+
const user = req.user.id;
3334

3435
const getIds = (data) => {
3536
userMessage = data.userMessage;
@@ -55,6 +56,7 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
5556
unfinished: true,
5657
cancelled: false,
5758
error: false,
59+
user,
5860
});
5961
}
6062

@@ -80,7 +82,7 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
8082
let response = await client.sendMessage(text, {
8183
getIds,
8284
// debug: true,
83-
user: req.user.id,
85+
user,
8486
conversationId,
8587
parentMessageId,
8688
overrideParentMessageId,
@@ -98,18 +100,18 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
98100
response.parentMessageId = overrideParentMessageId;
99101
}
100102

101-
await saveConvo(req.user.id, {
103+
await saveConvo(user, {
102104
...endpointOption,
103105
...endpointOption.modelOptions,
104106
conversationId,
105107
endpoint: 'anthropic',
106108
});
107109

108-
await saveMessage(response);
110+
await saveMessage({ ...response, user });
109111
sendMessage(res, {
110-
title: await getConvoTitle(req.user.id, conversationId),
112+
title: await getConvoTitle(user, conversationId),
111113
final: true,
112-
conversation: await getConvo(req.user.id, conversationId),
114+
conversation: await getConvo(user, conversationId),
113115
requestMessage: userMessage,
114116
responseMessage: response,
115117
});

‎api/server/routes/ask/askChatGPTBrowser.js

Lines changed: 11 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ router.post('/', setHeaders, async (req, res) => {
4848
});
4949

5050
if (!overrideParentMessageId) {
51-
await saveMessage(userMessage);
51+
await saveMessage({ ...userMessage, user: req.user.id });
5252
await saveConvo(req.user.id, {
5353
...userMessage,
5454
...endpointOption,
@@ -80,7 +80,7 @@ const ask = async ({
8080
res,
8181
}) => {
8282
let { text, parentMessageId: userParentMessageId, messageId: userMessageId } = userMessage;
83-
const userId = req.user.id;
83+
const user = req.user.id;
8484
let responseMessageId = crypto.randomUUID();
8585
let getPartialMessage = null;
8686
try {
@@ -100,6 +100,7 @@ const ask = async ({
100100
cancelled: false,
101101
error: false,
102102
isCreatedByUser: false,
103+
user,
103104
});
104105
}
105106
},
@@ -114,7 +115,7 @@ const ask = async ({
114115
conversationId,
115116
...endpointOption,
116117
abortController,
117-
userId,
118+
userId: user,
118119
onProgress: progressCallback.call(null, { res, text }),
119120
onEventMessage: (eventMessage) => {
120121
let data = null;
@@ -157,7 +158,7 @@ const ask = async ({
157158
isCreatedByUser: false,
158159
};
159160

160-
await saveMessage(responseMessage);
161+
await saveMessage({ ...responseMessage, user });
161162
responseMessage.messageId = newResponseMessageId;
162163

163164
// STEP2 update the conversation
@@ -181,7 +182,7 @@ const ask = async ({
181182
}
182183
}
183184

184-
await saveConvo(req.user.id, conversationUpdate);
185+
await saveConvo(user, conversationUpdate);
185186
conversationId = newConversationId;
186187

187188
// STEP3 update the user message
@@ -192,16 +193,17 @@ const ask = async ({
192193
if (!overrideParentMessageId) {
193194
await saveMessage({
194195
...userMessage,
196+
user,
195197
messageId: userMessageId,
196198
newMessageId: newUserMassageId,
197199
});
198200
}
199201
userMessageId = newUserMassageId;
200202

201203
sendMessage(res, {
202-
title: await getConvoTitle(req.user.id, conversationId),
204+
title: await getConvoTitle(user, conversationId),
203205
final: true,
204-
conversation: await getConvo(req.user.id, conversationId),
206+
conversation: await getConvo(user, conversationId),
205207
requestMessage: userMessage,
206208
responseMessage: responseMessage,
207209
});
@@ -210,7 +212,7 @@ const ask = async ({
210212
if (userParentMessageId == '00000000-0000-0000-0000-000000000000') {
211213
// const title = await titleConvo({ endpoint: endpointOption?.endpoint, text, response: responseMessage });
212214
const title = await response.details.title;
213-
await saveConvo(req.user.id, {
215+
await saveConvo(user, {
214216
conversationId: conversationId,
215217
title,
216218
});
@@ -227,7 +229,7 @@ const ask = async ({
227229
isCreatedByUser: false,
228230
text: `${getPartialMessage() ?? ''}\n\nError message: "${error.message}"`,
229231
};
230-
await saveMessage(errorMessage);
232+
await saveMessage({ ...errorMessage, user });
231233
handleError(res, errorMessage);
232234
}
233235
};

‎api/server/routes/ask/bingAI.js

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,7 @@ router.post('/', setHeaders, async (req, res) => {
6767
});
6868

6969
if (!overrideParentMessageId) {
70-
await saveMessage(userMessage);
70+
await saveMessage({ ...userMessage, user: req.user.id });
7171
await saveConvo(req.user.id, {
7272
...userMessage,
7373
...endpointOption,
@@ -100,6 +100,7 @@ const ask = async ({
100100
res,
101101
}) => {
102102
let { text, parentMessageId: userParentMessageId, messageId: userMessageId } = userMessage;
103+
const user = req.user.id;
103104

104105
let responseMessageId = crypto.randomUUID();
105106
const model = endpointOption?.jailbreak ? 'Sydney' : 'BingAI';
@@ -125,21 +126,22 @@ const ask = async ({
125126
cancelled: false,
126127
error: false,
127128
isCreatedByUser: false,
129+
user,
128130
});
129131
}
130132
},
131133
});
132134
const abortController = new AbortController();
133135
let bingConversationId = null;
134136
if (!isNewConversation) {
135-
const convo = await getConvo(req.user.id, conversationId);
137+
const convo = await getConvo(user, conversationId);
136138
bingConversationId = convo.bingConversationId;
137139
}
138140

139141
try {
140142
let response = await askBing({
141143
text,
142-
userId: req.user.id,
144+
userId: user,
143145
parentMessageId: userParentMessageId,
144146
conversationId: bingConversationId ?? conversationId,
145147
...endpointOption,
@@ -194,7 +196,7 @@ const ask = async ({
194196
isCreatedByUser: false,
195197
};
196198

197-
await saveMessage(responseMessage);
199+
await saveMessage({ ...responseMessage, user });
198200
responseMessage.messageId = newResponseMessageId;
199201

200202
let conversationUpdate = {
@@ -213,23 +215,24 @@ const ask = async ({
213215
conversationUpdate.invocationId = response.invocationId;
214216
}
215217

216-
await saveConvo(req.user.id, conversationUpdate);
218+
await saveConvo(user, conversationUpdate);
217219
userMessage.messageId = newUserMessageId;
218220

219221
// If response has parentMessageId, the fake userMessage.messageId should be updated to the real one.
220222
if (!overrideParentMessageId) {
221223
await saveMessage({
222224
...userMessage,
225+
user,
223226
messageId: userMessageId,
224227
newMessageId: newUserMessageId,
225228
});
226229
}
227230
userMessageId = newUserMessageId;
228231

229232
sendMessage(res, {
230-
title: await getConvoTitle(req.user.id, conversationId),
233+
title: await getConvoTitle(user, conversationId),
231234
final: true,
232-
conversation: await getConvo(req.user.id, conversationId),
235+
conversation: await getConvo(user, conversationId),
233236
requestMessage: userMessage,
234237
responseMessage: responseMessage,
235238
});
@@ -241,7 +244,7 @@ const ask = async ({
241244
response: responseMessage,
242245
});
243246

244-
await saveConvo(req.user.id, {
247+
await saveConvo(user, {
245248
conversationId: conversationId,
246249
title,
247250
});
@@ -263,12 +266,12 @@ const ask = async ({
263266
isCreatedByUser: false,
264267
};
265268

266-
saveMessage(responseMessage);
269+
saveMessage({ ...responseMessage, user });
267270

268271
return {
269-
title: await getConvoTitle(req.user.id, conversationId),
272+
title: await getConvoTitle(user, conversationId),
270273
final: true,
271-
conversation: await getConvo(req.user.id, conversationId),
274+
conversation: await getConvo(user, conversationId),
272275
requestMessage: userMessage,
273276
responseMessage: responseMessage,
274277
};
@@ -286,7 +289,7 @@ const ask = async ({
286289
model,
287290
isCreatedByUser: false,
288291
};
289-
await saveMessage(errorMessage);
292+
await saveMessage({ ...errorMessage, user });
290293
handleError(res, errorMessage);
291294
}
292295
}

0 commit comments

Comments
 (0)