anything-llm/server/models/workspace.js
Sean Hatfield 56fa17caf2
create configurable topN per workspace (#616)
* create configurable topN per workspace

* Update TopN UI text
Fix fallbacks for all providers
Add SQLite CHECK to TOPN value

* merge with master
Update zilliz provider for variable TopN

---------

Co-authored-by: timothycarambat <rambat1010@gmail.com>
2024-01-18 12:34:20 -08:00

213 lines
5.5 KiB
JavaScript

const prisma = require("../utils/prisma");
const slugify = require("slugify");
const { Document } = require("./documents");
const { WorkspaceUser } = require("./workspaceUsers");
const Workspace = {
writable: [
// Used for generic updates so we can validate keys in request body
"name",
"slug",
"vectorTag",
"openAiTemp",
"openAiHistory",
"lastUpdatedAt",
"openAiPrompt",
"similarityThreshold",
"chatModel",
"topN",
],
new: async function (name = null, creatorId = null) {
if (!name) return { result: null, message: "name cannot be null" };
var slug = slugify(name, { lower: true });
const existingBySlug = await this.get({ slug });
if (existingBySlug !== null) {
const slugSeed = Math.floor(10000000 + Math.random() * 90000000);
slug = slugify(`${name}-${slugSeed}`, { lower: true });
}
try {
const workspace = await prisma.workspaces.create({
data: { name, slug },
});
// If created with a user then we need to create the relationship as well.
// If creating with an admin User it wont change anything because admins can
// view all workspaces anyway.
if (!!creatorId) await WorkspaceUser.create(creatorId, workspace.id);
return { workspace, message: null };
} catch (error) {
console.error(error.message);
return { workspace: null, message: error.message };
}
},
update: async function (id = null, data = {}) {
if (!id) throw new Error("No workspace id provided for update");
const validKeys = Object.keys(data).filter((key) =>
this.writable.includes(key)
);
if (validKeys.length === 0)
return { workspace: { id }, message: "No valid fields to update!" };
try {
const workspace = await prisma.workspaces.update({
where: { id },
data,
});
return { workspace, message: null };
} catch (error) {
console.error(error.message);
return { workspace: null, message: error.message };
}
},
getWithUser: async function (user = null, clause = {}) {
if (["admin", "manager"].includes(user.role)) return this.get(clause);
try {
const workspace = await prisma.workspaces.findFirst({
where: {
...clause,
workspace_users: {
some: {
user_id: user?.id,
},
},
},
include: {
workspace_users: true,
documents: true,
},
});
if (!workspace) return null;
return {
...workspace,
documents: await Document.forWorkspace(workspace.id),
};
} catch (error) {
console.error(error.message);
return null;
}
},
get: async function (clause = {}) {
try {
const workspace = await prisma.workspaces.findFirst({
where: clause,
include: {
documents: true,
},
});
return workspace || null;
} catch (error) {
console.error(error.message);
return null;
}
},
delete: async function (clause = {}) {
try {
await prisma.workspaces.delete({
where: clause,
});
return true;
} catch (error) {
console.error(error.message);
return false;
}
},
where: async function (clause = {}, limit = null, orderBy = null) {
try {
const results = await prisma.workspaces.findMany({
where: clause,
...(limit !== null ? { take: limit } : {}),
...(orderBy !== null ? { orderBy } : {}),
});
return results;
} catch (error) {
console.error(error.message);
return [];
}
},
whereWithUser: async function (
user,
clause = {},
limit = null,
orderBy = null
) {
if (["admin", "manager"].includes(user.role))
return await this.where(clause, limit, orderBy);
try {
const workspaces = await prisma.workspaces.findMany({
where: {
...clause,
workspace_users: {
some: {
user_id: user.id,
},
},
},
...(limit !== null ? { take: limit } : {}),
...(orderBy !== null ? { orderBy } : {}),
});
return workspaces;
} catch (error) {
console.error(error.message);
return [];
}
},
whereWithUsers: async function (clause = {}, limit = null, orderBy = null) {
try {
const workspaces = await this.where(clause, limit, orderBy);
for (const workspace of workspaces) {
const userIds = (
await WorkspaceUser.where({ workspace_id: Number(workspace.id) })
).map((rel) => rel.user_id);
workspace.userIds = userIds;
}
return workspaces;
} catch (error) {
console.error(error.message);
return [];
}
},
updateUsers: async function (workspaceId, userIds = []) {
try {
await WorkspaceUser.delete({ workspace_id: Number(workspaceId) });
await WorkspaceUser.createManyUsers(userIds, workspaceId);
return { success: true, error: null };
} catch (error) {
console.error(error.message);
return { success: false, error: error.message };
}
},
resetWorkspaceChatModels: async () => {
try {
await prisma.workspaces.updateMany({
data: {
chatModel: null,
},
});
return { success: true, error: null };
} catch (error) {
console.error("Error resetting workspace chat models:", error.message);
return { success: false, error: error.message };
}
},
};
module.exports = { Workspace };