Merge pull request #8358 from amplication/feat/jovu-usage-limits

Feat/jovu usage limits

libs/util/billing-types/src/lib/billing-feature.types.ts

@@ -13,6 +13,7 @@ export enum BillingFeature {
   GitLab = "feature-gitlab",
   IgnoreValidationCodeGeneration = "feature-ignore-validation-code-generation",
   ImportDBSchema = "feature-import-db-schema",
+  JovuRequests = "feature-jovu-requests",
   Notification = "feature-notifications",
   Projects = "feature-projects",
   RedesignArchitecture = "feature-redesign-architecture",

packages/amplication-client/src/Assistant/Assistant.scss

@@ -76,6 +76,12 @@ $avatar-size: 30px;
     }
   }
 
+  &__limit {
+    padding: var(--default-spacing);
+    flex: 1;
+    text-align: center;
+  }
+
   &__header {
     display: flex;
     flex-direction: row;

packages/amplication-client/src/Assistant/Assistant.tsx

@@ -1,6 +1,12 @@
 import {
+  EnumContentAlign,
+  EnumFlexDirection,
+  EnumGapSize,
+  EnumItemsAlign,
+  EnumTextAlign,
   EnumTextColor,
   EnumTextStyle,
+  FlexItem,
   Icon,
   Text,
   TextField,
@@ -18,6 +24,10 @@ import classNames from "classnames";
 import { useAppContext } from "../context/appContext";
 import { Link } from "react-router-dom";
 import jovu from "../assets/jovu.svg";
+import { BillingFeature } from "@amplication/util-billing-types";
+import { useStiggContext } from "@stigg/react-sdk";
+import { GET_CONTACT_US_LINK } from "../Workspaces/queries/workspaceQueries";
+import { useQuery } from "@apollo/client";
 type SendMessageType = models.SendAssistantMessageInput;
 
 const INITIAL_VALUES: SendMessageType = {
@@ -53,6 +63,16 @@ const WIDTH_STATE_SETTINGS: Record<
 const Assistant = () => {
   const { currentWorkspace } = useAppContext();
 
+  const { stigg } = useStiggContext();
+
+  const { hasAccess } = stigg.getMeteredEntitlement({
+    featureId: BillingFeature.JovuRequests,
+  });
+
+  const { data } = useQuery(GET_CONTACT_US_LINK, {
+    variables: { id: currentWorkspace.id },
+  });
+
   const [open, setOpen] = useState(true);
   const [widthState, setWidthState] = useState(WIDTH_STATE_DEFAULT);
 
@@ -66,7 +86,6 @@ const Assistant = () => {
   const {
     sendMessage,
     messages,
-    sendMessageError: error,
     processingMessage: loading,
     streamError,
   } = useAssistant();
@@ -140,7 +159,51 @@ const Assistant = () => {
           </Tooltip>
         </div>
 
-        {currentWorkspace?.allowLLMFeatures ? (
+        {!hasAccess ? (
+          <FlexItem
+            direction={EnumFlexDirection.Column}
+            itemsAlign={EnumItemsAlign.Center}
+            contentAlign={EnumContentAlign.Center}
+            gap={EnumGapSize.Large}
+            className={`${CLASS_NAME}__limit`}
+          >
+            <img
+              src={jovu}
+              alt="jovu"
+              width={50}
+              height={50}
+              style={{ background: "white", borderRadius: "50%" }}
+            />
+            <Text textStyle={EnumTextStyle.H3} textAlign={EnumTextAlign.Center}>
+              You have reached the daily limit of Jovu requests for your plan.
+            </Text>
+            <Text
+              textStyle={EnumTextStyle.Tag}
+              textAlign={EnumTextAlign.Center}
+            >
+              Talk with us to upgrade and discover additional hidden
+              functionalities.
+            </Text>
+            <Text
+              textColor={EnumTextColor.White}
+              textStyle={EnumTextStyle.Tag}
+              textAlign={EnumTextAlign.Center}
+            >
+              <a
+                className={`${CLASS_NAME}__addon-section__contact-us`}
+                href={data?.contactUsLink}
+                target="blank"
+              >
+                <Text
+                  textColor={EnumTextColor.ThemeTurquoise}
+                  textStyle={EnumTextStyle.Tag}
+                >
+                  Talk with us
+                </Text>
+              </a>
+            </Text>
+          </FlexItem>
+        ) : currentWorkspace?.allowLLMFeatures ? (
           <>
             <div className={`${CLASS_NAME}__messages`}>
               {messages.map((message) => (
@@ -152,9 +215,6 @@ const Assistant = () => {
               ))}
 
               <div ref={messagesEndRef} />
-              {error && (
-                <div className={`${CLASS_NAME}__error`}>{error.message}</div>
-              )}
               {streamError && (
                 <div className={`${CLASS_NAME}__error`}>
                   {streamError.message}

packages/amplication-client/src/Assistant/hooks/useAssistant.tsx

@@ -50,18 +50,18 @@ const useAssistant = () => {
 
   const [threadId, setThreadId] = useState<string | null>(null);
 
-  const [
-    sendAssistantMessage,
-    { error: sendMessageError, loading: sendMessageLoading },
-  ] = useMutation<TAssistantThreadData>(SEND_ASSISTANT_MESSAGE, {
-    onCompleted: (data) => {
-      setThreadId(data.sendAssistantMessageWithStream.id);
-      setMessages([
-        ...messages,
-        ...data.sendAssistantMessageWithStream.messages,
-      ]);
-    },
-  });
+  const [sendAssistantMessage] = useMutation<TAssistantThreadData>(
+    SEND_ASSISTANT_MESSAGE,
+    {
+      onCompleted: (data) => {
+        setThreadId(data.sendAssistantMessageWithStream.id);
+        setMessages([
+          ...messages,
+          ...data.sendAssistantMessageWithStream.messages,
+        ]);
+      },
+    }
+  );
 
   const { error: streamError } = useSubscription<TAssistantMessageUpdatedData>(
     ASSISTANT_MESSAGE_UPDATED,
@@ -125,15 +125,29 @@ const useAssistant = () => {
         },
       },
     }).catch((error) => {
-      console.error(error);
+      setMessages([
+        ...messages,
+        {
+          text: "I'm sorry, I had a problem processing your request.",
+          role: models.EnumAssistantMessageRole.Assistant,
+          id: Date.now().toString() + "_",
+          createdAt: "",
+        },
+        {
+          text: error.message,
+          role: models.EnumAssistantMessageRole.Assistant,
+          id: Date.now().toString() + "_",
+          createdAt: "",
+        },
+      ]);
+
+      setProcessingMessage(false);
     });
   };
 
   return {
     sendMessage,
     messages,
-    sendMessageError,
-    sendMessageLoading,
     streamError,
     processingMessage,
   };

packages/amplication-server/src/core/assistant/assistant.module.ts

@@ -11,6 +11,7 @@ import { PluginCatalogModule } from "../pluginCatalog/pluginCatalog.module";
 import { PluginInstallationModule } from "../pluginInstallation/pluginInstallation.module";
 import { ModuleActionModule } from "../moduleAction/moduleAction.module";
 import { ModuleDtoModule } from "../moduleDto/moduleDto.module";
+import { BillingModule } from "../billing/billing.module";
 
 @Module({
   imports: [
@@ -23,6 +24,7 @@ import { ModuleDtoModule } from "../moduleDto/moduleDto.module";
     PluginInstallationModule,
     ModuleActionModule,
     ModuleDtoModule,
+    BillingModule,
   ],
   providers: [
     AssistantService,

packages/amplication-server/src/core/assistant/assistant.service.ts

@@ -32,6 +32,9 @@ import { EnumModuleActionGqlOperation } from "../moduleAction/dto/EnumModuleActi
 import { EnumModuleActionRestVerb } from "../moduleAction/dto/EnumModuleActionRestVerb";
 import { PropertyTypeDef } from "../moduleDto/dto/propertyTypes/PropertyTypeDef";
 import { EnumModuleActionRestInputSource } from "../moduleAction/dto/EnumModuleActionRestInputSource";
+import { BillingService } from "../billing/billing.service";
+import { BillingFeature } from "@amplication/util-billing-types";
+import { BillingLimitationError } from "../../errors/BillingLimitationError";
 
 enum EnumAssistantFunctions {
   CreateEntity = "createEntity",
@@ -88,6 +91,8 @@ export class AssistantService {
     private readonly pluginInstallationService: PluginInstallationService,
     private readonly moduleActionService: ModuleActionService,
     private readonly moduleDtoService: ModuleDtoService,
+    private readonly billingService: BillingService,
+
     configService: ConfigService
   ) {
     this.logger.info("starting assistant service");
@@ -140,11 +145,7 @@ export class AssistantService {
     };
   }
 
-  async processMessage(
-    messageText: string,
-    threadId: string,
-    context: AssistantContext
-  ): Promise<AssistantThread> {
+  async validateAndReportUsage(context: AssistantContext) {
     if (!this.assistantFeatureEnabled)
       throw new AmplicationError("The assistant AI feature is disabled");
 
@@ -154,6 +155,33 @@ export class AssistantService {
       );
     }
 
+    if (this.billingService.isBillingEnabled) {
+      const usage = await this.billingService.getMeteredEntitlement(
+        context.user.workspace.id,
+        BillingFeature.JovuRequests
+      );
+
+      if (usage && !usage?.hasAccess) {
+        throw new BillingLimitationError(
+          "This workspace has exceeded the number of allowed requests. Please upgrade your plan to continue using this feature.",
+          BillingFeature.JovuRequests
+        );
+      }
+
+      await this.billingService.reportUsage(
+        context.user.workspace.id,
+        BillingFeature.JovuRequests
+      );
+    }
+  }
+
+  async processMessage(
+    messageText: string,
+    threadId: string,
+    context: AssistantContext
+  ): Promise<AssistantThread> {
+    await this.validateAndReportUsage(context);
+
     const openai = this.openai;
 
     const preparedThread = await this.prepareThread(
@@ -186,14 +214,7 @@ export class AssistantService {
     threadId: string,
     context: AssistantContext
   ): Promise<AssistantThread> {
-    if (!this.assistantFeatureEnabled)
-      throw new AmplicationError("The assistant AI feature is disabled");
-
-    if (context.user.workspace.allowLLMFeatures === false) {
-      throw new AmplicationError(
-        "AI-powered features are disabled for this workspace"
-      );
-    }
+    await this.validateAndReportUsage(context);
 
     const openai = this.openai;