diff --git a/src/renderer/src/pages/agents/components/AddAgentPopup.tsx b/src/renderer/src/pages/agents/components/AddAgentPopup.tsx index 58671b2e..0d38e47d 100644 --- a/src/renderer/src/pages/agents/components/AddAgentPopup.tsx +++ b/src/renderer/src/pages/agents/components/AddAgentPopup.tsx @@ -8,14 +8,16 @@ import { useAgents } from '@renderer/hooks/useAgents' import { useSidebarIconShow } from '@renderer/hooks/useSidebarIcon' import { fetchGenerate } from '@renderer/services/ApiService' import { getDefaultModel } from '@renderer/services/AssistantService' +import { estimateTextTokens } from '@renderer/services/TokenService' import { useAppSelector } from '@renderer/store' import { Agent, KnowledgeBase } from '@renderer/types' import { getLeadingEmoji, uuid } from '@renderer/utils' import { Button, Form, FormInstance, Input, Modal, Popover, Select, SelectProps } from 'antd' import TextArea from 'antd/es/input/TextArea' -import { useRef, useState } from 'react' +import { useEffect, useRef, useState } from 'react' import { useTranslation } from 'react-i18next' import stringWidth from 'string-width' +import styled from 'styled-components' interface Props { resolve: (data: Agent | null) => void @@ -36,6 +38,7 @@ const PopupContainer: React.FC = ({ resolve }) => { const formRef = useRef(null) const [emoji, setEmoji] = useState('') const [loading, setLoading] = useState(false) + const [tokenCount, setTokenCount] = useState(0) const knowledgeState = useAppSelector((state) => state.knowledge) const showKnowledgeIcon = useSidebarIconShow('knowledge') const knowledgeOptions: SelectProps['options'] = [] @@ -47,6 +50,19 @@ const PopupContainer: React.FC = ({ resolve }) => { }) }) + useEffect(() => { + const updateTokenCount = async () => { + const prompt = formRef.current?.getFieldValue('prompt') + if (prompt) { + const count = await estimateTextTokens(prompt) + setTokenCount(count) + } else { + setTokenCount(0) + } + } + updateTokenCount() + }, [form.getFieldValue('prompt')]) + const onFinish = (values: FieldType) => { const _emoji = emoji || getLeadingEmoji(values.name) @@ -132,7 +148,13 @@ const PopupContainer: React.FC = ({ resolve }) => { labelAlign="left" colon={false} style={{ marginTop: 25 }} - onFinish={onFinish}> + onFinish={onFinish} + onValuesChange={async (changedValues) => { + if (changedValues.prompt) { + const count = await estimateTextTokens(changedValues.prompt) + setTokenCount(count) + } + }}> } arrow> @@ -147,7 +169,10 @@ const PopupContainer: React.FC = ({ resolve }) => { label={t('agents.add.prompt')} rules={[{ required: true }]} style={{ position: 'relative' }}> -