fix: show embedding token usage

2025-06-09 18:45:47 +08:00
parent c896c88c23
commit 1771ca363a
4 changed files with 128 additions and 92 deletions
@@ -5,14 +5,7 @@ import { CloseOutlined } from '@ant-design/icons';
 import { useIntl } from '@umijs/max';
 import { Button, Drawer } from 'antd';
 import _ from 'lodash';
-import React, {
-  memo,
-  useCallback,
-  useEffect,
-  useMemo,
-  useRef,
-  useState
-} from 'react';
+import React, { useEffect, useMemo, useRef, useState } from 'react';
 import styled from 'styled-components';
 import { queryCatalogItemSpec } from '../apis';
 import {
@@ -126,14 +119,22 @@ const AddModal: React.FC<AddModalProps> = (props) => {
  };

  // use for size change and quantization change
-  const pickSomeFieldsValue = () => {
+  const pickSomeFieldsValue = (defaultSpec: CatalogSpec) => {
    const formData = form.current?.getFieldsValue();
-    return _.pick(formData, [
+    const currentData = _.pick(formData, [
      'worker_selector',
      'gpu_selector',
      'env',
-      'backend_version'
+      'backend_version',
+      'backend_parameters'
    ]);
+    return {
+      ...currentData,
+      backend_parameters:
+        currentData.backend_parameters?.length > 0
+          ? currentData.backend_parameters
+          : defaultSpec.backend_parameters || []
+    };
  };

  const generateSubmitData = (formData: FormData) => {
@@ -431,7 +432,7 @@ const AddModal: React.FC<AddModalProps> = (props) => {
    });
    form.current.setFieldsValue({
      ...data,
-      ...pickSomeFieldsValue()
+      ...pickSomeFieldsValue(data)
    });
    handleCheckFormData();
  };
@@ -453,7 +454,7 @@ const AddModal: React.FC<AddModalProps> = (props) => {
    // set form data
    form.current.setFieldsValue({
      ...data,
-      ...pickSomeFieldsValue()
+      ...pickSomeFieldsValue(data)
    });
    handleCheckFormData();
  };
@@ -466,10 +467,10 @@ const AddModal: React.FC<AddModalProps> = (props) => {
    onOk(data);
  };

-  const handleCancel = useCallback(() => {
+  const handleCancel = () => {
    onCancel?.();
    axiosToken.current?.cancel?.();
-  }, [onCancel]);
+  };

  const showExtraButton = useMemo(() => {
    return warningStatus.show && warningStatus.type !== 'success';
@@ -614,4 +615,4 @@ const AddModal: React.FC<AddModalProps> = (props) => {
  );
 };

-export default memo(AddModal);
+export default AddModal;
@@ -38,6 +38,7 @@ import { generateEmbeddingCode } from '../view-code/embedding';
 import DynamicParams from './dynamic-params';
 import FileList from './file-list';
 import InputList from './input-list';
+import TokenUsage from './token-usage';
 import ViewCommonCode from './view-common-code';

 interface MessageProps {
@@ -525,21 +526,29 @@ const GroundEmbedding: React.FC<MessageProps> = forwardRef((props, ref) => {
                onSelect={handleonSelect}
                onPaste={handleOnPaste}
              ></InputList>
-              <div style={{ marginTop: 8 }}>
-                <FileList
-                  fileList={fileList}
-                  textListCount={textList.length || 0}
-                  onDelete={handleDeleteFile}
-                ></FileList>
-              </div>
-              {lessTwoInput && (
-                <AlertInfo
-                  type="danger"
-                  message={intl.formatMessage({
-                    id: 'playground.documents.verify.embedding'
-                  })}
-                ></AlertInfo>
+              {fileList.length > 0 && (
+                <div style={{ marginTop: 8 }}>
+                  <FileList
+                    fileList={fileList}
+                    textListCount={textList.length || 0}
+                    onDelete={handleDeleteFile}
+                  ></FileList>
+                </div>
              )}
+              {lessTwoInput && (
+                <div className="m-t-16">
+                  <AlertInfo
+                    type="danger"
+                    message={intl.formatMessage({
+                      id: 'playground.documents.verify.embedding'
+                    })}
+                  ></AlertInfo>
+                </div>
+              )}
+              <TokenUsage
+                tokenResult={tokenResult}
+                className="m-t-16"
+              ></TokenUsage>
            </div>
          </div>
        </div>
@@ -33,6 +33,7 @@ import '../style/system-message-wrap.less';
 import { generateRerankCode } from '../view-code/rerank';
 import DynamicParams from './dynamic-params';
 import InputList from './input-list';
+import TokenUsage from './token-usage';
 import ViewCommonCode from './view-common-code';

 interface MessageProps {
@@ -179,7 +180,7 @@ const GroundReranker: React.FC<MessageProps> = forwardRef((props, ref) => {
    return res * 100;
  };

-  const renderPercent = useCallback((data: any) => {
+  const renderPercent = (data: any) => {
    if (!data.showExtra || !data.percent) {
      return null;
    }
@@ -208,7 +209,7 @@ const GroundReranker: React.FC<MessageProps> = forwardRef((props, ref) => {
        </span>
      </div>
    );
-  }, []);
+  };

  const setMessageId = () => {
    messageId.current = messageId.current + 1;
@@ -320,59 +321,52 @@ const GroundReranker: React.FC<MessageProps> = forwardRef((props, ref) => {
    inputListRef.current?.handleAdd();
  };

-  const handleTextListChange = useCallback(
-    (list: { text: string; uid: number | string; name: string }[]) => {
-      const newList = list?.map((item: any) => {
-        item.percent = undefined;
-        item.score = undefined;
-        item.rank = undefined;
-        return item;
+  const handleTextListChange = (
+    list: { text: string; uid: number | string; name: string }[]
+  ) => {
+    const newList = list?.map((item: any) => {
+      item.percent = undefined;
+      item.score = undefined;
+      item.rank = undefined;
+      return item;
+    });
+    setTextList(newList);
+  };
+
+  const handleonSelect = (data: {
+    start: number;
+    end: number;
+    beforeText: string;
+    afterText: string;
+    index: number;
+  }) => {
+    selectionTextRef.current = data;
+  };
+
+  const handleOnPaste = (e: any, index: number) => {
+    if (!multiplePasteEnable.current) return;
+    const text = e.clipboardData.getData('text');
+    if (text) {
+      const dataLlist = text.split('\n').map((item: string) => {
+        return {
+          text: item?.trim(),
+          name: '',
+          uid: setMessageId(),
+          percent: undefined,
+          score: undefined,
+          rank: undefined
+        };
      });
-      setTextList(newList);
-    },
-    []
-  );
+      dataLlist[0].text = `${selectionTextRef.current?.beforeText || ''}${dataLlist[0].text}${selectionTextRef.current?.afterText || ''}`;
+      const result = [
+        ...textList.slice(0, index),
+        ...dataLlist,
+        ...textList.slice(index + 1)
+      ].filter((item) => item.text);

-  const handleonSelect = useCallback(
-    (data: {
-      start: number;
-      end: number;
-      beforeText: string;
-      afterText: string;
-      index: number;
-    }) => {
-      selectionTextRef.current = data;
-    },
-    []
-  );
-
-  const handleOnPaste = useCallback(
-    (e: any, index: number) => {
-      if (!multiplePasteEnable.current) return;
-      const text = e.clipboardData.getData('text');
-      if (text) {
-        const dataLlist = text.split('\n').map((item: string) => {
-          return {
-            text: item?.trim(),
-            name: '',
-            uid: setMessageId(),
-            percent: undefined,
-            score: undefined,
-            rank: undefined
-          };
-        });
-        dataLlist[0].text = `${selectionTextRef.current?.beforeText || ''}${dataLlist[0].text}${selectionTextRef.current?.afterText || ''}`;
-        const result = [
-          ...textList.slice(0, index),
-          ...dataLlist,
-          ...textList.slice(index + 1)
-        ].filter((item) => item.text);
-
-        setTextList(result);
-      }
-    },
-    [textList]
-  );
+      setTextList(result);
+    }
+  };

  const renderExtra = useMemo(() => {
    if (modelMeta?.n_ctx && modelMeta?.n_slot) {
@@ -469,15 +463,7 @@ const GroundReranker: React.FC<MessageProps> = forwardRef((props, ref) => {
                  {intl.formatMessage({ id: 'playground.embedding.documents' })}
                </span>
              </h3>
-              <span className="m-l-10 font-size-12">
-                {' '}
-                {tokenResult?.total_tokens && (
-                  <span style={{ color: 'var(--ant-orange)' }}>
-                    {intl.formatMessage({ id: 'playground.tokenusage' })}:{' '}
-                    {tokenResult?.total_tokens}
-                  </span>
-                )}
-              </span>
+
              <div className="flex-center gap-10">
                <Tooltip
                  title={intl.formatMessage({
@@ -532,6 +518,10 @@ const GroundReranker: React.FC<MessageProps> = forwardRef((props, ref) => {
                  ></AlertInfo>
                </div>
              )}
+              <TokenUsage
+                tokenResult={tokenResult}
+                className="m-t-16"
+              ></TokenUsage>
            </div>
          </div>
          <div></div>
@@ -0,0 +1,36 @@
+import { useIntl } from '@umijs/max';
+import React from 'react';
+import styled from 'styled-components';
+
+const TokenUsageWrapper = styled.div`
+  font-size: var(--font-size-small);
+  display: flex;
+  align-items: center;
+  justify-content: center;
+  padding: 2px 5px;
+  .text {
+    color: var(--ant-orange);
+  }
+`;
+
+const TokenUsage: React.FC<{
+  tokenResult?: any;
+  [key: string]: any;
+}> = ({ tokenResult, ...rest }) => {
+  const intl = useIntl();
+  if (!tokenResult) {
+    return null;
+  }
+  return (
+    <TokenUsageWrapper {...rest}>
+      {tokenResult?.total_tokens && (
+        <span className="text">
+          {intl.formatMessage({ id: 'playground.tokenusage' })}:{' '}
+          {tokenResult?.total_tokens}
+        </span>
+      )}
+    </TokenUsageWrapper>
+  );
+};
+
+export default TokenUsage;