ui: add support for response limiting

2025-05-21 15:24:29 +02:00 · 2025-05-21 15:24:29 +02:00 · 2ad4c8732d
commit 2ad4c8732d
parent 20bb27c67c
8 changed files with 203 additions and 70 deletions
--- a/frontend/app/components/Kai/KaiChat.tsx
+++ b/frontend/app/components/Kai/KaiChat.tsx
@ -98,7 +98,9 @@ function KaiChat() {
  };
  return (
    <div className="w-full mx-auto" style={{ maxWidth: PANEL_SIZES.maxWidth }}>
-      <div className={'w-full rounded-lg overflow-hidden border shadow'}>
+      <div
+        className={'w-full rounded-lg overflow-hidden border shadow relative'}
+      >
        <ChatHeader
          chatTitle={chatTitle}
          openChats={openChats}
--- a/frontend/app/components/Kai/KaiService.ts
+++ b/frontend/app/components/Kai/KaiService.ts
@ -3,7 +3,7 @@ import AiService from '@/services/AiService';
 export default class KaiService extends AiService {
  getKaiChats = async (
    projectId: string,
-  ): Promise<{ title: string; thread_id: string, datetime: string }[]> => {
+  ): Promise<{ title: string; thread_id: string; datetime: string }[]> => {
    const r = await this.client.get(`/kai/${projectId}/chats`);
    if (!r.ok) {
      throw new Error('Failed to fetch chats');
@ -81,24 +81,45 @@ export default class KaiService extends AiService {
    return data;
  };

-  getMsgChart = async (messageId: string, projectId: string): Promise<{ filters: any[], chart: string, eventsOrder: string }> => {
-    const r = await this.client.get(`/kai/${projectId}/chats/data/${messageId}`);
+  getMsgChart = async (
+    messageId: string,
+    projectId: string,
+  ): Promise<{ filters: any[]; chart: string; eventsOrder: string }> => {
+    const r = await this.client.get(
+      `/kai/${projectId}/chats/data/${messageId}`,
+    );
    if (!r.ok) {
      throw new Error('Failed to fetch chart data');
    }
    const data = await r.json();
    return data;
-  }
+  };

-  saveChartData = async (messageId: string, projectId: string, chartData: any) => {
-    const r = await this.client.post(`/kai/${projectId}/chats/data/${messageId}`, {
-      chart_data: JSON.stringify(chartData),
-    });
+  saveChartData = async (
+    messageId: string,
+    projectId: string,
+    chartData: any,
+  ) => {
+    const r = await this.client.post(
+      `/kai/${projectId}/chats/data/${messageId}`,
+      {
+        chart_data: JSON.stringify(chartData),
+      },
+    );
    if (!r.ok) {
      throw new Error('Failed to save chart data');
    }

    const data = await r.json();
    return data;
-  }
+  };
+
+  checkUsage = async (): Promise<{ total: number; used: number }> => {
+    const r = await this.client.get(`/kai/usage`);
+    if (!r.ok) {
+      throw new Error('Failed to fetch usage');
+    }
+    const data = await r.json();
+    return data;
+  };
 }
--- a/frontend/app/components/Kai/KaiStore.ts
+++ b/frontend/app/components/Kai/KaiStore.ts
@ -31,9 +31,15 @@ class KaiStore {
  queryText = '';
  loadingChat = false;
  replacing: string | null = null;
+  usage = {
+    total: 0,
+    used: 0,
+    percent: 0,
+  };

  constructor() {
    makeAutoObservable(this);
+    this.checkUsage();
  }

  get lastHumanMessage() {
@ -146,6 +152,7 @@ class KaiStore {
      console.error('No token found');
      return;
    }
+    this.checkUsage();
    this.chatManager = new ChatManager({ ...settings, token });
    this.chatManager.setOnMsgHook({
      msgCallback: (msg) => {
@ -184,6 +191,7 @@ class KaiStore {
              supports_visualization: msg.supports_visualization,
              chart_data: '',
            };
+            this.bumpUsage();
            this.addMessage(msgObj);
            this.setProcessingStage(null);
          }
@ -201,6 +209,11 @@ class KaiStore {
    this.replacing = replacing;
  };

+  bumpUsage = () => {
+    this.usage.used += 1;
+    this.usage.percent = (this.usage.used / this.usage.total) * 100;
+  };
+
  sendMessage = (message: string) => {
    if (this.chatManager) {
      this.chatManager.sendMessage(message, !!this.replacing);
@ -329,6 +342,19 @@ class KaiStore {
    const parsedData = JSON.parse(data);
    return new Widget().fromJson(parsedData);
  };
+
+  checkUsage = async () => {
+    try {
+      const { total, used } = await kaiService.checkUsage();
+      this.usage = {
+        total,
+        used,
+        percent: (used / total) * 100,
+      };
+    } catch (e) {
+      console.error(e);
+    }
+  };
 }

 export const kaiStore = new KaiStore();
--- a/frontend/app/components/Kai/components/ChatHeader.tsx
+++ b/frontend/app/components/Kai/components/ChatHeader.tsx
@ -1,6 +1,7 @@
 import React from 'react';
 import { Icon } from 'UI';
 import { MessagesSquare, ArrowLeft } from 'lucide-react';
+import { useTranslation } from 'react-i18next';

 function ChatHeader({
  openChats = () => {},
@ -11,6 +12,7 @@ function ChatHeader({
  openChats?: () => void;
  chatTitle: string | null;
 }) {
+  const { t } = useTranslation();
  return (
    <div
      className={
@ -20,17 +22,21 @@ function ChatHeader({
      <div className={'flex-1'}>
        {goBack ? (
          <div
-            className={'flex items-center gap-2 font-semibold cursor-pointer'}
+            className={
+              'w-fit flex items-center gap-2 font-semibold cursor-pointer'
+            }
            onClick={goBack}
          >
            <ArrowLeft size={14} />
-            <div>Back</div>
+            <div>{t('Back')}</div>
          </div>
        ) : null}
      </div>
      <div className={'flex items-center gap-2 mx-auto max-w-[80%]'}>
        {chatTitle ? (
-          <div className="font-semibold text-xl whitespace-nowrap truncate">{chatTitle}</div>
+          <div className="font-semibold text-xl whitespace-nowrap truncate">
+            {chatTitle}
+          </div>
        ) : (
          <>
            <Icon name={'kai_colored'} size={18} />
@ -38,14 +44,14 @@ function ChatHeader({
          </>
        )}
      </div>
-      <div
-        className={
-          'font-semibold cursor-pointer flex items-center gap-2 flex-1 justify-end'
-        }
-        onClick={openChats}
-      >
-        <MessagesSquare size={14} />
-        <div>Chats</div>
+      <div className={'flex-1 justify-end flex items-center gap-2'}>
+        <div
+          className="font-semibold w-fit cursor-pointer flex items-center gap-2"
+          onClick={openChats}
+        >
+          <MessagesSquare size={14} />
+          <div>{t('Chats')}</div>
+        </div>
      </div>
    </div>
  );
--- a/frontend/app/components/Kai/components/ChatInput.tsx
+++ b/frontend/app/components/Kai/components/ChatInput.tsx
@ -3,6 +3,7 @@ import { Button, Input, Tooltip } from 'antd';
 import { SendHorizonal, OctagonX } from 'lucide-react';
 import { kaiStore } from '../KaiStore';
 import { observer } from 'mobx-react-lite';
+import Usage from './Usage';

 function ChatInput({
  isLoading,
@ -13,7 +14,9 @@ function ChatInput({
  onSubmit: (str: string) => void;
  threadId: string;
 }) {
-  const inputRef = React.useRef<Input>(null);
+  const inputRef = React.useRef<typeof Input>(null);
+  const usage = kaiStore.usage;
+  const limited = usage.percent >= 100;
  const inputValue = kaiStore.queryText;
  const isProcessing = kaiStore.processingStage !== null;
  const setInputValue = (text: string) => {
@ -21,6 +24,9 @@ function ChatInput({
  };

  const submit = () => {
+    if (limited) {
+      return;
+    }
    if (isProcessing) {
      const settings = { projectId: '2325', userId: '0', threadId };
      void kaiStore.cancelGeneration(settings);
@ -46,50 +52,57 @@ function ChatInput({
  const isReplacing = kaiStore.replacing !== null;

  return (
-    <Input
-      onPressEnter={submit}
-      onKeyDown={(e) => {
-        if (e.key === 'Escape') {
-          cancelReplace();
-        }
-      }}
-      ref={inputRef}
-      placeholder={'Ask anything about your product and users...'}
-      size={'large'}
-      value={inputValue}
-      onChange={(e) => setInputValue(e.target.value)}
-      suffix={
-        <>
-          {isReplacing ? (
-            <Tooltip title={'Cancel replacing'}>
+    <div className="relative">
+      <Input
+        onPressEnter={submit}
+        onKeyDown={(e) => {
+          if (e.key === 'Escape') {
+            cancelReplace();
+          }
+        }}
+        ref={inputRef}
+        placeholder={'Ask anything about your product and users...'}
+        size={'large'}
+        value={inputValue}
+        onChange={(e) => setInputValue(e.target.value)}
+        suffix={
+          <>
+            {isReplacing ? (
+              <Tooltip title={'Cancel Editing'}>
+                <Button
+                  onClick={cancelReplace}
+                  icon={<OctagonX size={16} />}
+                  type={'text'}
+                  size={'small'}
+                  shape={'circle'}
+                  disabled={limited}
+                />
+              </Tooltip>
+            ) : null}
+            <Tooltip title={'Send message'}>
              <Button
-                onClick={cancelReplace}
-                icon={<OctagonX size={16} />}
+                loading={isLoading}
+                onClick={submit}
+                disabled={limited}
+                icon={
+                  isProcessing ? (
+                    <OctagonX size={16} />
+                  ) : (
+                    <SendHorizonal size={16} />
+                  )
+                }
                type={'text'}
                size={'small'}
                shape={'circle'}
              />
            </Tooltip>
-          ) : null}
-          <Tooltip title={'Send message'}>
-            <Button
-              loading={isLoading}
-              onClick={submit}
-              icon={
-                isProcessing ? (
-                  <OctagonX size={16} />
-                ) : (
-                  <SendHorizonal size={16} />
-                )
-              }
-              type={'text'}
-              size={'small'}
-              shape={'circle'}
-            />
-          </Tooltip>
-        </>
-      }
-    />
+          </>
+        }
+      />
+      <div className="absolute ml-1 top-2 -right-11">
+        <Usage />
+      </div>
+    </div>
  );
 }

--- a/frontend/app/components/Kai/components/ChatMsg.tsx
+++ b/frontend/app/components/Kai/components/ChatMsg.tsx
@ -19,7 +19,7 @@ import { toast } from 'react-toastify';
 import { durationFormatted } from 'App/date';
 import WidgetChart from '@/components/Dashboard/components/WidgetChart';
 import Widget from 'App/mstore/types/widget';
-import cn from 'classnames';
+import { useTranslation } from 'react-i18next';

 function ChatMsg({
  userName,
@ -32,6 +32,7 @@ function ChatMsg({
  canEdit?: boolean;
  siteId: string;
 }) {
+  const { t } = useTranslation();
  const [metric, setMetric] = React.useState<Widget | null>(null);
  const [loadingChart, setLoadingChart] = React.useState(false);
  const {
@ -46,9 +47,13 @@ function ChatMsg({
  const isEditing = kaiStore.replacing && messageId === kaiStore.replacing;
  const [isProcessing, setIsProcessing] = React.useState(false);
  const bodyRef = React.useRef<HTMLDivElement>(null);
-  const onRetry = () => {
+  const onEdit = () => {
    kaiStore.editMessage(text, messageId);
  };
+  const onCancelEdit = () => {
+    kaiStore.setQueryText('');
+    kaiStore.setReplacing(null);
+  }
  const onFeedback = (feedback: 'like' | 'dislike', messageId: string) => {
    kaiStore.sendMsgFeedback(feedback, messageId, siteId);
  };
@ -143,17 +148,31 @@ function ChatMsg({
          </div>
        ) : null}
        {isUser ? (
-          canEdit ? (
+          <>
            <div
-              onClick={onRetry}
-              className={
-                'ml-auto flex items-center gap-2 px-2 rounded-lg border border-gray-medium text-sm cursor-pointer hover:border-main hover:text-main w-fit'
-              }
+              onClick={onEdit}
+              className={cn(
+                'ml-auto flex items-center gap-2 px-2',
+                'rounded-lg border border-gray-medium text-sm cursor-pointer',
+                'hover:border-main hover:text-main w-fit',
+                canEdit && !isEditing ? '' : 'hidden',
+              )}
            >
              <ListRestart size={16} />
-              <div>Edit</div>
+              <div>{t('Edit')}</div>
            </div>
-          ) : null
+            <div
+              onClick={onCancelEdit}
+              className={cn(
+                'ml-auto flex items-center gap-2 px-2',
+                'rounded-lg border border-gray-medium text-sm cursor-pointer',
+                'hover:border-main hover:text-main w-fit',
+                isEditing ? '' : 'hidden',
+              )}
+            >
+              <div>{t('Cancel')}</div>
+            </div>
+          </>
        ) : (
          <div className={'flex items-center gap-2'}>
            {duration ? <MsgDuration duration={duration} /> : null}
--- a/frontend/app/components/Kai/components/ChatsModal.tsx
+++ b/frontend/app/components/Kai/components/ChatsModal.tsx
@ -5,6 +5,8 @@ import { MessagesSquare, Trash } from 'lucide-react';
 import { kaiService } from 'App/services';
 import { toast } from 'react-toastify';
 import { useTranslation } from 'react-i18next';
+import { kaiStore } from '../KaiStore';
+import { observer } from 'mobx-react-lite';

 function ChatsModal({
  onSelect,
@ -14,6 +16,7 @@ function ChatsModal({
  projectId: string;
 }) {
  const { t } = useTranslation();
+  const { usage } = kaiStore;
  const {
    data = [],
    isPending,
@ -24,6 +27,10 @@ function ChatsModal({
    staleTime: 1000 * 60,
  });

+  React.useEffect(() => {
+    kaiStore.checkUsage();
+  }, []);
+
  const datedCollections = React.useMemo(() => {
    return data.length ? splitByDate(data) : [];
  }, [data.length]);
@ -42,6 +49,14 @@ function ChatsModal({
        <MessagesSquare size={16} />
        <span>{t('Chats')}</span>
      </div>
+      {usage.percent > 80 ? (
+        <div className="text-red text-sm">
+          {t('You have used {{used}} out of {{total}} daily requests', {
+            used: usage.used,
+            total: usage.total,
+          })}
+        </div>
+      ) : null}
      {isPending ? (
        <div className="animate-pulse text-disabled-text">{t('Loading chats')}...</div>
      ) : (
@ -118,4 +133,4 @@ function ChatsList({
  );
 }

-export default ChatsModal;
+export default observer(ChatsModal);
--- a/frontend/app/components/Kai/components/Usage.tsx
+++ b/frontend/app/components/Kai/components/Usage.tsx
@ -0,0 +1,31 @@
+import React from 'react';
+import { kaiStore } from '../KaiStore';
+import { observer } from 'mobx-react-lite';
+import { Progress, Tooltip } from 'antd';
+const getUsageColor = (percent: number) => {
+  return 'disabled-text';
+};
+
+function Usage() {
+  const { usage } = kaiStore;
+  const color = getUsageColor(usage.percent);
+
+  if (usage.total === 0) {
+    return null;
+  }
+  return (
+    <div>
+      <Tooltip title={`Daily response limit (${usage.used}/${usage.total})`}>
+        <Progress
+          percent={usage.percent}
+          strokeColor={usage.percent < 99 ? 'var(--color-main)' : 'var(--color-red)'}
+          showInfo={false}
+          type="circle"
+          size={24}
+        />
+      </Tooltip>
+    </div>
+  );
+}
+
+export default observer(Usage);