anything-llm/server/models/workspaceChats.js
Sean Hatfield 8b5d9ccdb3
[FEAT] Fork chat to new thread (#1788)
* implement thread forking feature

* rename thread based on forked message

* refactor bulk message create for thread fork + bump prisma version

* revert prisma version bump

* add todo to bulkCreate function in workspace chats

* cast user input to expected type to prevent prisma injection

* refactor: update order of ops for thread fork

---------

Co-authored-by: Timothy Carambat <rambat1010@gmail.com>
2024-07-03 14:44:35 -07:00

263 lines
6.5 KiB
JavaScript

const prisma = require("../utils/prisma");
const WorkspaceChats = {
new: async function ({
workspaceId,
prompt,
response = {},
user = null,
threadId = null,
include = true,
}) {
try {
const chat = await prisma.workspace_chats.create({
data: {
workspaceId,
prompt,
response: JSON.stringify(response),
user_id: user?.id || null,
thread_id: threadId,
include,
},
});
return { chat, message: null };
} catch (error) {
console.error(error.message);
return { chat: null, message: error.message };
}
},
forWorkspaceByUser: async function (
workspaceId = null,
userId = null,
limit = null,
orderBy = null
) {
if (!workspaceId || !userId) return [];
try {
const chats = await prisma.workspace_chats.findMany({
where: {
workspaceId,
user_id: userId,
thread_id: null, // this function is now only used for the default thread on workspaces and users
include: true,
},
...(limit !== null ? { take: limit } : {}),
...(orderBy !== null ? { orderBy } : { orderBy: { id: "asc" } }),
});
return chats;
} catch (error) {
console.error(error.message);
return [];
}
},
forWorkspace: async function (
workspaceId = null,
limit = null,
orderBy = null
) {
if (!workspaceId) return [];
try {
const chats = await prisma.workspace_chats.findMany({
where: {
workspaceId,
thread_id: null, // this function is now only used for the default thread on workspaces
include: true,
},
...(limit !== null ? { take: limit } : {}),
...(orderBy !== null ? { orderBy } : { orderBy: { id: "asc" } }),
});
return chats;
} catch (error) {
console.error(error.message);
return [];
}
},
markHistoryInvalid: async function (workspaceId = null, user = null) {
if (!workspaceId) return;
try {
await prisma.workspace_chats.updateMany({
where: {
workspaceId,
user_id: user?.id,
thread_id: null, // this function is now only used for the default thread on workspaces
},
data: {
include: false,
},
});
return;
} catch (error) {
console.error(error.message);
}
},
markThreadHistoryInvalid: async function (
workspaceId = null,
user = null,
threadId = null
) {
if (!workspaceId || !threadId) return;
try {
await prisma.workspace_chats.updateMany({
where: {
workspaceId,
thread_id: threadId,
user_id: user?.id,
},
data: {
include: false,
},
});
return;
} catch (error) {
console.error(error.message);
}
},
get: async function (clause = {}, limit = null, orderBy = null) {
try {
const chat = await prisma.workspace_chats.findFirst({
where: clause,
...(limit !== null ? { take: limit } : {}),
...(orderBy !== null ? { orderBy } : {}),
});
return chat || null;
} catch (error) {
console.error(error.message);
return null;
}
},
delete: async function (clause = {}) {
try {
await prisma.workspace_chats.deleteMany({
where: clause,
});
return true;
} catch (error) {
console.error(error.message);
return false;
}
},
where: async function (
clause = {},
limit = null,
orderBy = null,
offset = null
) {
try {
const chats = await prisma.workspace_chats.findMany({
where: clause,
...(limit !== null ? { take: limit } : {}),
...(offset !== null ? { skip: offset } : {}),
...(orderBy !== null ? { orderBy } : {}),
});
return chats;
} catch (error) {
console.error(error.message);
return [];
}
},
count: async function (clause = {}) {
try {
const count = await prisma.workspace_chats.count({
where: clause,
});
return count;
} catch (error) {
console.error(error.message);
return 0;
}
},
whereWithData: async function (
clause = {},
limit = null,
offset = null,
orderBy = null
) {
const { Workspace } = require("./workspace");
const { User } = require("./user");
try {
const results = await this.where(clause, limit, orderBy, offset);
for (const res of results) {
const workspace = await Workspace.get({ id: res.workspaceId });
res.workspace = workspace
? { name: workspace.name, slug: workspace.slug }
: { name: "deleted workspace", slug: null };
const user = res.user_id ? await User.get({ id: res.user_id }) : null;
res.user = user
? { username: user.username }
: { username: "unknown user" };
}
return results;
} catch (error) {
console.error(error.message);
return [];
}
},
updateFeedbackScore: async function (chatId = null, feedbackScore = null) {
if (!chatId) return;
try {
await prisma.workspace_chats.update({
where: {
id: Number(chatId),
},
data: {
feedbackScore:
feedbackScore === null ? null : Number(feedbackScore) === 1,
},
});
return;
} catch (error) {
console.error(error.message);
}
},
// Explicit update of settings + key validations.
// Only use this method when directly setting a key value
// that takes no user input for the keys being modified.
_update: async function (id = null, data = {}) {
if (!id) throw new Error("No workspace chat id provided for update");
try {
await prisma.workspace_chats.update({
where: { id },
data,
});
return true;
} catch (error) {
console.error(error.message);
return false;
}
},
bulkCreate: async function (chatsData) {
// TODO: Replace with createMany when we update prisma to latest version
// The version of prisma that we are currently using does not support createMany with SQLite
try {
const createdChats = [];
for (const chatData of chatsData) {
const chat = await prisma.workspace_chats.create({
data: chatData,
});
createdChats.push(chat);
}
return { chats: createdChats, message: null };
} catch (error) {
console.error(error.message);
return { chats: null, message: error.message };
}
},
};
module.exports = { WorkspaceChats };