From e658c1c2cf2826488765910e2ea60070bcb893dd Mon Sep 17 00:00:00 2001 From: Stanislas Polu Date: Tue, 16 Jan 2024 18:32:12 +0100 Subject: [PATCH] Fix "runner: bump limits for paid users + refactor paid plan conditions"" (#3246) * Revert "Revert "runner: bump limits for paid users + refactor paid plan conditions"" * fix import * more * lint --- .../GalleryAssistantPreviewContainer.tsx | 4 +- .../assistant_builder/AssistantBuilder.tsx | 13 +-- front/lib/api/assistant/agent.ts | 8 +- front/lib/api/assistant/generation.ts | 6 +- front/lib/api/assistant/global_agents.ts | 104 +++++++----------- front/lib/api/assistant/pubsub.ts | 3 + front/lib/auth.ts | 5 + front/lib/plans/plan_codes.ts | 14 +++ front/pages/api/w/[wId]/data_sources/index.ts | 3 +- front/pages/poke/[wId]/index.tsx | 7 +- front/pages/w/[wId]/members/index.tsx | 8 +- front/pages/w/[wId]/subscription/index.tsx | 5 +- 12 files changed, 84 insertions(+), 96 deletions(-) diff --git a/front/components/assistant/GalleryAssistantPreviewContainer.tsx b/front/components/assistant/GalleryAssistantPreviewContainer.tsx index 67dd2c9a92e7..27f1fd2af1f4 100644 --- a/front/components/assistant/GalleryAssistantPreviewContainer.tsx +++ b/front/components/assistant/GalleryAssistantPreviewContainer.tsx @@ -12,7 +12,7 @@ import { SendNotificationsContext, } from "@app/components/sparkle/Notification"; import { isLargeModel } from "@app/lib/assistant"; -import { FREE_TEST_PLAN_CODE } from "@app/lib/plans/plan_codes"; +import { isUpgraded } from "@app/lib/plans/plan_codes"; import { PostAgentListStatusRequestBody } from "@app/pages/api/w/[wId]/members/me/agent_list_status"; type AssistantPreviewFlow = "personal" | "workspace"; @@ -176,7 +176,7 @@ export function GalleryAssistantPreviewContainer({ const isGlobal = scope === "global"; const isAddedToWorkspace = flow === "workspace" && isAdded; - const hasAccessToLargeModels = plan?.code !== FREE_TEST_PLAN_CODE; + const hasAccessToLargeModels = isUpgraded(plan); const eligibleForTesting = hasAccessToLargeModels || !isLargeModel(generation?.model); const isTestable = !isGlobal && !isAdded && eligibleForTesting; diff --git a/front/components/assistant_builder/AssistantBuilder.tsx b/front/components/assistant_builder/AssistantBuilder.tsx index f7eb244cf3ad..35b67b612c0c 100644 --- a/front/components/assistant_builder/AssistantBuilder.tsx +++ b/front/components/assistant_builder/AssistantBuilder.tsx @@ -67,7 +67,7 @@ import { SendNotificationsContext } from "@app/components/sparkle/Notification"; import { getSupportedModelConfig } from "@app/lib/assistant"; import { CONNECTOR_CONFIGURATIONS } from "@app/lib/connector_providers"; import { isActivatedStructuredDB } from "@app/lib/development"; -import { FREE_TEST_PLAN_CODE } from "@app/lib/plans/plan_codes"; +import { isUpgraded } from "@app/lib/plans/plan_codes"; import { useSlackChannelsLinkedWithAgent } from "@app/lib/swr"; import { classNames } from "@app/lib/utils"; @@ -307,10 +307,9 @@ export default function AssistantBuilder({ scope: defaultScope, generationSettings: { ...DEFAULT_ASSISTANT_STATE.generationSettings, - modelSettings: - plan.code === FREE_TEST_PLAN_CODE - ? GPT_3_5_TURBO_MODEL_CONFIG - : GPT_4_TURBO_MODEL_CONFIG, + modelSettings: !isUpgraded(plan) + ? GPT_3_5_TURBO_MODEL_CONFIG + : GPT_4_TURBO_MODEL_CONFIG, }, } ); @@ -1531,9 +1530,7 @@ function AdvancedSettings({ {usedModelConfigs - .filter( - (m) => !(m.largeModel && plan.code === FREE_TEST_PLAN_CODE) - ) + .filter((m) => !(m.largeModel && !isUpgraded(plan))) .map((modelConfig) => ( { const owner = auth.workspace(); - const plan = auth.plan(); - if (!owner || !plan) { + if (!owner) { throw new Error("Unexpected unauthenticated call to `runGeneration`"); } @@ -304,7 +302,7 @@ export async function* runGeneration( let model = c.model; - if (isLargeModel(model) && plan.code === FREE_TEST_PLAN_CODE) { + if (isLargeModel(model) && !auth.isUpgraded()) { yield { type: "generation_error", created: Date.now(), diff --git a/front/lib/api/assistant/global_agents.ts b/front/lib/api/assistant/global_agents.ts index 9f8b41a3b40e..01b9f80d4e31 100644 --- a/front/lib/api/assistant/global_agents.ts +++ b/front/lib/api/assistant/global_agents.ts @@ -18,12 +18,11 @@ import { MISTRAL_MEDIUM_MODEL_CONFIG, MISTRAL_SMALL_MODEL_CONFIG, } from "@dust-tt/types"; -import { DustAPI, GlobalAgentStatus, PlanType } from "@dust-tt/types"; +import { DustAPI, GlobalAgentStatus } from "@dust-tt/types"; import { GLOBAL_AGENTS_SID } from "@app/lib/assistant"; import { Authenticator, prodAPICredentialsForOwner } from "@app/lib/auth"; import { GlobalAgentSettings } from "@app/lib/models/assistant/agent"; -import { FREE_TEST_PLAN_CODE } from "@app/lib/plans/plan_codes"; import logger from "@app/logger/logger"; class HelperAssistantPrompt { @@ -82,20 +81,15 @@ async function _getHelperGlobalAgent( if (!owner) { throw new Error("Unexpected `auth` without `workspace`."); } - const plan = auth.plan(); - if (!plan) { - throw new Error("Unexpected `auth` without `plan`."); - } - const model = - plan.code === FREE_TEST_PLAN_CODE - ? { - providerId: GPT_3_5_TURBO_MODEL_CONFIG.providerId, - modelId: GPT_3_5_TURBO_MODEL_CONFIG.modelId, - } - : { - providerId: GPT_4_TURBO_MODEL_CONFIG.providerId, - modelId: GPT_4_TURBO_MODEL_CONFIG.modelId, - }; + const model = !auth.isUpgraded() + ? { + providerId: GPT_3_5_TURBO_MODEL_CONFIG.providerId, + modelId: GPT_3_5_TURBO_MODEL_CONFIG.modelId, + } + : { + providerId: GPT_4_TURBO_MODEL_CONFIG.providerId, + modelId: GPT_4_TURBO_MODEL_CONFIG.modelId, + }; return { id: -1, sId: GLOBAL_AGENTS_SID.HELPER, @@ -151,12 +145,11 @@ async function _getGPT35TurboGlobalAgent({ } async function _getGPT4GlobalAgent({ - plan, + auth, }: { - plan: PlanType; + auth: Authenticator; }): Promise { - const status = - plan.code === FREE_TEST_PLAN_CODE ? "disabled_free_workspace" : "active"; + const status = !auth.isUpgraded() ? "disabled_free_workspace" : "active"; return { id: -1, sId: GLOBAL_AGENTS_SID.GPT4, @@ -216,14 +209,13 @@ async function _getClaudeInstantGlobalAgent({ } async function _getClaudeGlobalAgent({ + auth, settings, - plan, }: { + auth: Authenticator; settings: GlobalAgentSettings | null; - plan: PlanType; }): Promise { - const status = - plan.code === FREE_TEST_PLAN_CODE ? "disabled_free_workspace" : "active"; + const status = !auth.isUpgraded() ? "disabled_free_workspace" : "active"; return { id: -1, sId: GLOBAL_AGENTS_SID.CLAUDE, @@ -250,14 +242,14 @@ async function _getClaudeGlobalAgent({ } async function _getMistralMediumGlobalAgent({ - plan, + auth, settings, }: { - plan: PlanType; + auth: Authenticator; settings: GlobalAgentSettings | null; }): Promise { let status = settings?.status ?? "disabled_by_admin"; - if (plan.code === FREE_TEST_PLAN_CODE) { + if (!auth.isUpgraded()) { status = "disabled_free_workspace"; } @@ -376,11 +368,6 @@ async function _getManagedDataSourceAgent( throw new Error("Unexpected `auth` without `workspace`."); } - const plan = auth.plan(); - if (!plan) { - throw new Error("Unexpected `auth` without `plan`."); - } - const prodCredentials = await prodAPICredentialsForOwner(owner); // Check if deactivated by an admin @@ -439,16 +426,15 @@ async function _getManagedDataSourceAgent( generation: { id: -1, prompt, - model: - plan.code === FREE_TEST_PLAN_CODE - ? { - providerId: GPT_3_5_TURBO_MODEL_CONFIG.providerId, - modelId: GPT_3_5_TURBO_MODEL_CONFIG.modelId, - } - : { - providerId: GPT_4_TURBO_MODEL_CONFIG.providerId, - modelId: GPT_4_TURBO_MODEL_CONFIG.modelId, - }, + model: !auth.isUpgraded() + ? { + providerId: GPT_3_5_TURBO_MODEL_CONFIG.providerId, + modelId: GPT_3_5_TURBO_MODEL_CONFIG.modelId, + } + : { + providerId: GPT_4_TURBO_MODEL_CONFIG.providerId, + modelId: GPT_4_TURBO_MODEL_CONFIG.modelId, + }, temperature: 0.4, }, action: { @@ -565,10 +551,8 @@ async function _getNotionGlobalAgent( async function _getDustGlobalAgent( auth: Authenticator, { - plan, settings, }: { - plan: PlanType; settings: GlobalAgentSettings | null; } ): Promise { @@ -645,16 +629,15 @@ async function _getDustGlobalAgent( id: -1, prompt: "Assist the user based on the retrieved data from their workspace. Unlesss the user explicitely asks for a detailed answer, you goal is to provide a quick answer to their question.", - model: - plan.code === FREE_TEST_PLAN_CODE - ? { - providerId: GPT_3_5_TURBO_MODEL_CONFIG.providerId, - modelId: GPT_3_5_TURBO_MODEL_CONFIG.modelId, - } - : { - providerId: GPT_4_TURBO_MODEL_CONFIG.providerId, - modelId: GPT_4_TURBO_MODEL_CONFIG.modelId, - }, + model: !auth.isUpgraded() + ? { + providerId: GPT_3_5_TURBO_MODEL_CONFIG.providerId, + modelId: GPT_3_5_TURBO_MODEL_CONFIG.modelId, + } + : { + providerId: GPT_4_TURBO_MODEL_CONFIG.providerId, + modelId: GPT_4_TURBO_MODEL_CONFIG.modelId, + }, temperature: 0.4, }, action: { @@ -691,11 +674,6 @@ export async function getGlobalAgent( throw new Error("Cannot find Global Agent Configuration: no workspace."); } - const plan = auth.plan(); - if (!plan) { - throw new Error("Unexpected `auth` without `plan`."); - } - if (preFetchedDataSources === null) { const prodCredentials = await prodAPICredentialsForOwner(owner); const api = new DustAPI(prodCredentials, logger); @@ -719,18 +697,18 @@ export async function getGlobalAgent( agentConfiguration = await _getGPT35TurboGlobalAgent({ settings }); break; case GLOBAL_AGENTS_SID.GPT4: - agentConfiguration = await _getGPT4GlobalAgent({ plan }); + agentConfiguration = await _getGPT4GlobalAgent({ auth }); break; case GLOBAL_AGENTS_SID.CLAUDE_INSTANT: agentConfiguration = await _getClaudeInstantGlobalAgent({ settings }); break; case GLOBAL_AGENTS_SID.CLAUDE: - agentConfiguration = await _getClaudeGlobalAgent({ settings, plan }); + agentConfiguration = await _getClaudeGlobalAgent({ auth, settings }); break; case GLOBAL_AGENTS_SID.MISTRAL_MEDIUM: agentConfiguration = await _getMistralMediumGlobalAgent({ - plan, settings, + auth, }); break; case GLOBAL_AGENTS_SID.MISTRAL_SMALL: @@ -764,7 +742,7 @@ export async function getGlobalAgent( }); break; case GLOBAL_AGENTS_SID.DUST: - agentConfiguration = await _getDustGlobalAgent(auth, { plan, settings }); + agentConfiguration = await _getDustGlobalAgent(auth, { settings }); break; default: return null; diff --git a/front/lib/api/assistant/pubsub.ts b/front/lib/api/assistant/pubsub.ts index 56f6470be4c1..0ce3ac3074ab 100644 --- a/front/lib/api/assistant/pubsub.ts +++ b/front/lib/api/assistant/pubsub.ts @@ -55,6 +55,9 @@ export async function postUserMessageWithPubSub( let rateLimitKey: string | undefined = ""; if (auth.user()?.id) { maxPerTimeframe = 50; + if (auth.isUpgraded()) { + maxPerTimeframe = 200; + } timeframeSeconds = 60 * 60 * 3; rateLimitKey = `postUserMessageUser:${auth.user()?.id}`; } else { diff --git a/front/lib/auth.ts b/front/lib/auth.ts index 470c63b8ae58..371490d8bce0 100644 --- a/front/lib/auth.ts +++ b/front/lib/auth.ts @@ -20,6 +20,7 @@ import { Workspace, } from "@app/lib/models"; import { FREE_TEST_PLAN_DATA, PlanAttributes } from "@app/lib/plans/free_plans"; +import { isUpgraded } from "@app/lib/plans/plan_codes"; import { new_id } from "@app/lib/utils"; import logger from "@app/logger/logger"; import { authOptions } from "@app/pages/api/auth/[...nextauth]"; @@ -326,6 +327,10 @@ export class Authenticator { return this._subscription ? this._subscription.plan : null; } + isUpgraded(): boolean { + return isUpgraded(this.plan()); + } + /** * This is a convenience method to get the user from the Authenticator. The returned UserType * object won't have the user's workspaces set. diff --git a/front/lib/plans/plan_codes.ts b/front/lib/plans/plan_codes.ts index 6f62d0b1e043..4597322a38a0 100644 --- a/front/lib/plans/plan_codes.ts +++ b/front/lib/plans/plan_codes.ts @@ -1,3 +1,5 @@ +import { PlanType } from "@dust-tt/types"; + // Current free plans: export const FREE_UPGRADED_PLAN_CODE = "FREE_UPGRADED_PLAN"; export const FREE_TEST_PLAN_CODE = "FREE_TEST_PLAN"; @@ -9,3 +11,15 @@ export const PRO_PLAN_SEAT_29_CODE = "PRO_PLAN_SEAT_29"; * ENT_PLAN_FAKE is not subscribable and is only used to display the Enterprise plan in the UI (hence it's not stored on the db). */ export const ENT_PLAN_FAKE_CODE = "ENT_PLAN_FAKE_CODE"; + +/** + * `isUpgraded` returns true if the plan has access to all features of Dust, including large + * language models (meaning it's either a paid plan or free plan with (eg friends and family, or + * free trial plan)). + * + * Note: We didn't go for isFree or isPayingWorkspace as we have "upgraded" plans that are free. + */ +export const isUpgraded = (plan: PlanType | null): boolean => { + if (!plan) return false; + return plan.code !== FREE_TEST_PLAN_CODE; +}; diff --git a/front/pages/api/w/[wId]/data_sources/index.ts b/front/pages/api/w/[wId]/data_sources/index.ts index f4f59897db2d..60961398585d 100644 --- a/front/pages/api/w/[wId]/data_sources/index.ts +++ b/front/pages/api/w/[wId]/data_sources/index.ts @@ -7,7 +7,6 @@ import { NextApiRequest, NextApiResponse } from "next"; import { getDataSources } from "@app/lib/api/data_sources"; import { Authenticator, getSession } from "@app/lib/auth"; import { DataSource } from "@app/lib/models"; -import { FREE_TEST_PLAN_CODE } from "@app/lib/plans/plan_codes"; import logger from "@app/logger/logger"; import { apiError, withLogging } from "@app/logger/withlogging"; @@ -140,7 +139,7 @@ async function handler( splitter_id: "base_v0", max_chunk_size: dataSourceMaxChunkSize, qdrant_config: - plan.code !== FREE_TEST_PLAN_CODE && NODE_ENV === "production" + auth.isUpgraded() && NODE_ENV === "production" ? { cluster: "dedicated-1", shadow_write_cluster: null, diff --git a/front/pages/poke/[wId]/index.tsx b/front/pages/poke/[wId]/index.tsx index a6d89b9ce453..025b1c2cc9da 100644 --- a/front/pages/poke/[wId]/index.tsx +++ b/front/pages/poke/[wId]/index.tsx @@ -32,6 +32,7 @@ import { useSubmitFunction } from "@app/lib/client/utils"; import { FREE_TEST_PLAN_CODE, FREE_UPGRADED_PLAN_CODE, + isUpgraded, } from "@app/lib/plans/plan_codes"; import { getPlanInvitation } from "@app/lib/plans/subscription"; import { usePokePlans } from "@app/lib/swr"; @@ -651,7 +652,7 @@ const WorkspacePage = ({ variant="secondaryWarning" onClick={onDowngrade} disabled={ - subscription.plan.code === FREE_TEST_PLAN_CODE || + !isUpgraded(subscription.plan) || workspaceHasManagedDataSources } /> @@ -661,9 +662,7 @@ const WorkspacePage = ({ label="Upgrade to free upgraded plan" variant="tertiary" onClick={onUpgrade} - disabled={ - subscription.plan.code !== FREE_TEST_PLAN_CODE - } + disabled={isUpgraded(subscription.plan)} /> diff --git a/front/pages/w/[wId]/members/index.tsx b/front/pages/w/[wId]/members/index.tsx index 153130072ad2..9ecb76a074f6 100644 --- a/front/pages/w/[wId]/members/index.tsx +++ b/front/pages/w/[wId]/members/index.tsx @@ -28,7 +28,7 @@ import AppLayout from "@app/components/sparkle/AppLayout"; import { subNavigationAdmin } from "@app/components/sparkle/navigation"; import { SendNotificationsContext } from "@app/components/sparkle/Notification"; import { Authenticator, getSession, getUserFromSession } from "@app/lib/auth"; -import { FREE_TEST_PLAN_CODE } from "@app/lib/plans/plan_codes"; +import { isUpgraded } from "@app/lib/plans/plan_codes"; import { useMembers, useWorkspaceInvitations } from "@app/lib/swr"; import { classNames, isEmailValid } from "@app/lib/utils"; @@ -153,8 +153,7 @@ export default function WorkspaceAdmin({ size="sm" icon={Cog6ToothIcon} onClick={() => { - if (plan.code === FREE_TEST_PLAN_CODE) - setShowNoInviteLinkPopup(true); + if (!isUpgraded(plan)) setShowNoInviteLinkPopup(true); else setInviteSettingsModalOpen(true); }} /> @@ -278,8 +277,7 @@ export default function WorkspaceAdmin({ size="sm" icon={PlusIcon} onClick={() => { - if (plan.code === FREE_TEST_PLAN_CODE) - setShowNoInviteFreePlanPopup(true); + if (!isUpgraded(plan)) setShowNoInviteFreePlanPopup(true); else if (subscription.paymentFailingSince) setShowNoInviteFailedPaymentPopup(true); else setInviteEmailModalOpen(true); diff --git a/front/pages/w/[wId]/subscription/index.tsx b/front/pages/w/[wId]/subscription/index.tsx index 9c4237b0ce54..2a61180dffaf 100644 --- a/front/pages/w/[wId]/subscription/index.tsx +++ b/front/pages/w/[wId]/subscription/index.tsx @@ -23,6 +23,7 @@ import { useSubmitFunction } from "@app/lib/client/utils"; import { FREE_TEST_PLAN_CODE, FREE_UPGRADED_PLAN_CODE, + isUpgraded, PRO_PLAN_SEAT_29_CODE, } from "@app/lib/plans/plan_codes"; import { getPlanInvitation } from "@app/lib/plans/subscription"; @@ -131,7 +132,7 @@ export default function Subscription({ if (content.checkoutUrl) { await router.push(content.checkoutUrl); } else if (content.success) { - await router.reload(); // We cannot swr the plan so we just reload the page. + router.reload(); // We cannot swr the plan so we just reload the page. } } }); @@ -166,7 +167,7 @@ export default function Subscription({ const isProcessing = isSubscribingPlan || isGoingToStripePortal; const plan = subscription.plan; - const chipColor = plan.code === FREE_TEST_PLAN_CODE ? "emerald" : "sky"; + const chipColor = !isUpgraded(plan) ? "emerald" : "sky"; const onClickProPlan = async () => handleSubscribePlan(); const onClickEnterprisePlan = () => {