Skip to content

Commit

Permalink
Enforce 1MB limit even for upgraded workspaces (#1894)
Browse files Browse the repository at this point in the history
* Tweak poke

* Enforce 1MB even for upgraded workspaces

* migration to enforce 1mb on upgraded workspaces

* nit alert
  • Loading branch information
spolu authored Sep 29, 2023
1 parent 4d6d903 commit 63112f1
Show file tree
Hide file tree
Showing 6 changed files with 93 additions and 23 deletions.
3 changes: 2 additions & 1 deletion front/lib/api/workspace.ts
Original file line number Diff line number Diff line change
Expand Up @@ -127,7 +127,8 @@ export async function upgradeWorkspace(workspaceId: number) {

plan.limits.dataSources.count = -1;
plan.limits.dataSources.documents.count = -1;
plan.limits.dataSources.documents.sizeMb = -1;
// Enforce 1MB limitation even for upgraded workspaces.
plan.limits.dataSources.documents.sizeMb = 1;
plan.limits.dataSources.managed = true;
plan.limits.largeModels = true;

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,58 @@
import { Op } from "sequelize";

import { planForWorkspace } from "@app/lib/auth";
import { Workspace } from "@app/lib/models";

const { LIVE = false } = process.env;

async function main() {
console.log("Fetching Upgraded Worspaces...");
const workspaces = await Workspace.findAll({
where: {
upgradedAt: {
[Op.not]: null,
},
},
});
console.log(
`Found ${workspaces.length} workspaces for which to add largeModels = true`
);

const chunkSize = 16;
const chunks = [];
for (let i = 0; i < workspaces.length; i += chunkSize) {
chunks.push(workspaces.slice(i, i + chunkSize));
}

for (let i = 0; i < chunks.length; i++) {
console.log(`Processing chunk ${i}/${chunks.length}...`);
const chunk = chunks[i];
await Promise.all(
chunk.map((workspace: Workspace) => {
return set1MBLimit(!!LIVE, workspace);
})
);
}
}

async function set1MBLimit(live: boolean, workspace: Workspace) {
const plan = planForWorkspace(workspace);
plan.limits.dataSources.documents.sizeMb = 1;
if (live) {
await workspace.update({
plan: JSON.stringify(plan),
});
} else {
console.log(`Would have mutated ${workspace.sId}`);
}
}

main()
.then(() => {
console.log("done");
process.exit(0);
})
.catch((err) => {
console.error(err);
process.exit(1);
});
Original file line number Diff line number Diff line change
Expand Up @@ -270,7 +270,7 @@ async function handler(
api_error: {
type: "data_source_quota_error",
message:
"Data sources document upload size is limited to 1MB on our free plan. Contact [email protected] if you want to increase it.",
"Data sources document upload size is limited to 1MB. Contact [email protected] if you want to increase it.",
},
});
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -187,7 +187,7 @@ async function handler(
api_error: {
type: "data_source_quota_error",
message:
"Data sources document upload size is limited to 1MB on our free plan. Contact [email protected] if you want to increase it.",
"Data sources document upload size is limited to 1MB. Contact [email protected] if you want to increase it.",
},
});
}
Expand Down
18 changes: 6 additions & 12 deletions front/pages/poke/[wId]/index.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -248,14 +248,6 @@ const WorkspacePage = ({
(ds) => !!ds.connectorProvider
);

// @todo remove fallback on workspace plan
const isFullyUpgraded =
workspace.upgradedAt !== null ||
(workspace.plan?.limits.dataSources.count === -1 &&
workspace.plan?.limits.dataSources.documents.count === -1 &&
workspace.plan?.limits.dataSources.documents.sizeMb === -1 &&
workspace.plan?.limits.dataSources.managed === true);

return (
<div className="min-h-screen bg-structure-50">
<PokeNavbar />
Expand All @@ -275,7 +267,7 @@ const WorkspacePage = ({
<div className="flex justify-center">
<div className="mx-2 w-1/3">
<h2 className="text-md mb-4 font-bold">Plan:</h2>
{isFullyUpgraded ? (
{workspace.upgradedAt ? (
<p className="mb-4 text-green-600">
This workspace was fully upgraded
{workspace.upgradedAt &&
Expand All @@ -294,18 +286,20 @@ const WorkspacePage = ({
label="Downgrade"
variant="secondaryWarning"
onClick={onDowngrade}
disabled={!isFullyUpgraded || workspaceHasManagedDataSources}
disabled={
!workspace.upgradedAt || workspaceHasManagedDataSources
}
/>

<Button
label="Upgrade"
variant="secondary"
onClick={onUpgrade}
disabled={isFullyUpgraded}
disabled={!!workspace.upgradedAt}
/>
</div>
</div>
{isFullyUpgraded && workspaceHasManagedDataSources && (
{workspace.upgradedAt && workspaceHasManagedDataSources && (
<span className="mx-2 w-1/3">
<p className="text-warning mb-4 text-sm ">
Delete managed data sources before downgrading.
Expand Down
33 changes: 25 additions & 8 deletions front/pages/w/[wId]/builder/data-sources/[name]/upsert.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -116,10 +116,27 @@ export default function DataSourceUpsert({
}
}, [dataSource.name, loadDocumentId, owner.sId]);

const alertDataSourcesLimit = () => {
window.alert(
"DataSource document upload size is limited to " +
`${owner.plan.limits.dataSources.documents.sizeMb}MB (of raw text)` +
". Contact [email protected] if you want to increase this limit."
);
};

const handleFileLoadedText = (e: any) => {
const content = e.target.result;
setText(content);
setUploading(false);

// Enforce plan limits: DataSource documents size.
if (
owner.plan.limits.dataSources.documents.sizeMb != -1 &&
text.length > 1024 * 1024 * owner.plan.limits.dataSources.documents.sizeMb
) {
alertDataSourcesLimit();
return;
}
setText(content);
};

const handleFileLoadedPDF = async (e: any) => {
Expand All @@ -133,21 +150,21 @@ export default function DataSourceUpsert({
const strings = content.items.map((item: any) => item.str);
text += strings.join(" ") + "\n";
}
setText(text);

setUploading(false);
};

const handleFileUpload = async (file: File) => {
// Enforce plan limits: DataSource documents size.
if (
owner.plan.limits.dataSources.documents.sizeMb != -1 &&
file.size > 1024 * 1024 * owner.plan.limits.dataSources.documents.sizeMb
text.length > 1024 * 1024 * owner.plan.limits.dataSources.documents.sizeMb
) {
window.alert(
"DataSource document upload size is limited to 1MB on our free plan. Contact [email protected] if you want to increase this limit."
);
alertDataSourcesLimit();
return;
}
setText(text);
};

const handleFileUpload = async (file: File) => {
setUploading(true);
if (file.type === "application/pdf") {
const fileReader = new FileReader();
Expand Down

0 comments on commit 63112f1

Please sign in to comment.