feat(AddAgentPopup, AssistantPromptSettings): Add token count display for prompts

- Implement token counting functionality in both AddAgentPopup and AssistantPromptSettings components.
- Display the token count dynamically as the user types in the prompt text area.
- Refactor text area components to include a styled token count indicator.
This commit is contained in:
George·Dong 2025-03-12 23:57:19 +08:00 committed by 亢奋猫
parent 8b2c1cbe99
commit 8a3bf652d3
2 changed files with 85 additions and 13 deletions

View File

@ -8,14 +8,16 @@ import { useAgents } from '@renderer/hooks/useAgents'
import { useSidebarIconShow } from '@renderer/hooks/useSidebarIcon' import { useSidebarIconShow } from '@renderer/hooks/useSidebarIcon'
import { fetchGenerate } from '@renderer/services/ApiService' import { fetchGenerate } from '@renderer/services/ApiService'
import { getDefaultModel } from '@renderer/services/AssistantService' import { getDefaultModel } from '@renderer/services/AssistantService'
import { estimateTextTokens } from '@renderer/services/TokenService'
import { useAppSelector } from '@renderer/store' import { useAppSelector } from '@renderer/store'
import { Agent, KnowledgeBase } from '@renderer/types' import { Agent, KnowledgeBase } from '@renderer/types'
import { getLeadingEmoji, uuid } from '@renderer/utils' import { getLeadingEmoji, uuid } from '@renderer/utils'
import { Button, Form, FormInstance, Input, Modal, Popover, Select, SelectProps } from 'antd' import { Button, Form, FormInstance, Input, Modal, Popover, Select, SelectProps } from 'antd'
import TextArea from 'antd/es/input/TextArea' import TextArea from 'antd/es/input/TextArea'
import { useRef, useState } from 'react' import { useEffect, useRef, useState } from 'react'
import { useTranslation } from 'react-i18next' import { useTranslation } from 'react-i18next'
import stringWidth from 'string-width' import stringWidth from 'string-width'
import styled from 'styled-components'
interface Props { interface Props {
resolve: (data: Agent | null) => void resolve: (data: Agent | null) => void
@ -36,6 +38,7 @@ const PopupContainer: React.FC<Props> = ({ resolve }) => {
const formRef = useRef<FormInstance>(null) const formRef = useRef<FormInstance>(null)
const [emoji, setEmoji] = useState('') const [emoji, setEmoji] = useState('')
const [loading, setLoading] = useState(false) const [loading, setLoading] = useState(false)
const [tokenCount, setTokenCount] = useState(0)
const knowledgeState = useAppSelector((state) => state.knowledge) const knowledgeState = useAppSelector((state) => state.knowledge)
const showKnowledgeIcon = useSidebarIconShow('knowledge') const showKnowledgeIcon = useSidebarIconShow('knowledge')
const knowledgeOptions: SelectProps['options'] = [] const knowledgeOptions: SelectProps['options'] = []
@ -47,6 +50,19 @@ const PopupContainer: React.FC<Props> = ({ 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 onFinish = (values: FieldType) => {
const _emoji = emoji || getLeadingEmoji(values.name) const _emoji = emoji || getLeadingEmoji(values.name)
@ -132,7 +148,13 @@ const PopupContainer: React.FC<Props> = ({ resolve }) => {
labelAlign="left" labelAlign="left"
colon={false} colon={false}
style={{ marginTop: 25 }} style={{ marginTop: 25 }}
onFinish={onFinish}> onFinish={onFinish}
onValuesChange={async (changedValues) => {
if (changedValues.prompt) {
const count = await estimateTextTokens(changedValues.prompt)
setTokenCount(count)
}
}}>
<Form.Item name="name" label="Emoji"> <Form.Item name="name" label="Emoji">
<Popover content={<EmojiPicker onEmojiClick={setEmoji} />} arrow> <Popover content={<EmojiPicker onEmojiClick={setEmoji} />} arrow>
<Button icon={emoji && <span style={{ fontSize: 20 }}>{emoji}</span>}>{t('common.select')}</Button> <Button icon={emoji && <span style={{ fontSize: 20 }}>{emoji}</span>}>{t('common.select')}</Button>
@ -147,7 +169,10 @@ const PopupContainer: React.FC<Props> = ({ resolve }) => {
label={t('agents.add.prompt')} label={t('agents.add.prompt')}
rules={[{ required: true }]} rules={[{ required: true }]}
style={{ position: 'relative' }}> style={{ position: 'relative' }}>
<TextArea placeholder={t('agents.add.prompt.placeholder')} spellCheck={false} rows={10} /> <TextAreaContainer>
<TextArea placeholder={t('agents.add.prompt.placeholder')} spellCheck={false} rows={10} />
<TokenCount>Tokens: {tokenCount}</TokenCount>
</TextAreaContainer>
</Form.Item> </Form.Item>
<Button <Button
icon={loading ? <LoadingOutlined /> : <ThunderboltOutlined />} icon={loading ? <LoadingOutlined /> : <ThunderboltOutlined />}
@ -177,6 +202,23 @@ const PopupContainer: React.FC<Props> = ({ resolve }) => {
) )
} }
const TextAreaContainer = styled.div`
position: relative;
width: 100%;
`
const TokenCount = styled.div`
position: absolute;
bottom: 8px;
right: 8px;
background-color: var(--color-background-soft);
padding: 2px 8px;
border-radius: 4px;
font-size: 12px;
color: var(--color-text-2);
user-select: none;
`
export default class AddAgentPopup { export default class AddAgentPopup {
static topviewId = 0 static topviewId = 0
static hide() { static hide() {

View File

@ -3,11 +3,12 @@ import 'emoji-picker-element'
import { CloseCircleFilled } from '@ant-design/icons' import { CloseCircleFilled } from '@ant-design/icons'
import EmojiPicker from '@renderer/components/EmojiPicker' import EmojiPicker from '@renderer/components/EmojiPicker'
import { Box, HStack } from '@renderer/components/Layout' import { Box, HStack } from '@renderer/components/Layout'
import { estimateTextTokens } from '@renderer/services/TokenService'
import { Assistant, AssistantSettings } from '@renderer/types' import { Assistant, AssistantSettings } from '@renderer/types'
import { getLeadingEmoji } from '@renderer/utils' import { getLeadingEmoji } from '@renderer/utils'
import { Button, Input, Popover } from 'antd' import { Button, Input, Popover } from 'antd'
import TextArea from 'antd/es/input/TextArea' import TextArea from 'antd/es/input/TextArea'
import { useState } from 'react' import { useEffect, useState } from 'react'
import { useTranslation } from 'react-i18next' import { useTranslation } from 'react-i18next'
import styled from 'styled-components' import styled from 'styled-components'
@ -22,8 +23,17 @@ const AssistantPromptSettings: React.FC<Props> = ({ assistant, updateAssistant,
const [emoji, setEmoji] = useState(getLeadingEmoji(assistant.name) || assistant.emoji) const [emoji, setEmoji] = useState(getLeadingEmoji(assistant.name) || assistant.emoji)
const [name, setName] = useState(assistant.name.replace(getLeadingEmoji(assistant.name) || '', '').trim()) const [name, setName] = useState(assistant.name.replace(getLeadingEmoji(assistant.name) || '', '').trim())
const [prompt, setPrompt] = useState(assistant.prompt) const [prompt, setPrompt] = useState(assistant.prompt)
const [tokenCount, setTokenCount] = useState(0)
const { t } = useTranslation() const { t } = useTranslation()
useEffect(() => {
const updateTokenCount = async () => {
const count = await estimateTextTokens(prompt)
setTokenCount(count)
}
updateTokenCount()
}, [prompt])
const onUpdate = () => { const onUpdate = () => {
const _assistant = { ...assistant, name: name.trim(), emoji, prompt } const _assistant = { ...assistant, name: name.trim(), emoji, prompt }
updateAssistant(_assistant) updateAssistant(_assistant)
@ -81,15 +91,18 @@ const AssistantPromptSettings: React.FC<Props> = ({ assistant, updateAssistant,
<Box mt={8} mb={8} style={{ fontWeight: 'bold' }}> <Box mt={8} mb={8} style={{ fontWeight: 'bold' }}>
{t('common.prompt')} {t('common.prompt')}
</Box> </Box>
<TextArea <TextAreaContainer>
rows={10} <TextArea
placeholder={t('common.assistant') + t('common.prompt')} rows={10}
value={prompt} placeholder={t('common.assistant') + t('common.prompt')}
onChange={(e) => setPrompt(e.target.value)} value={prompt}
onBlur={onUpdate} onChange={(e) => setPrompt(e.target.value)}
spellCheck={false} onBlur={onUpdate}
style={{ minHeight: 'calc(80vh - 200px)', maxHeight: 'calc(80vh - 150px)' }} spellCheck={false}
/> style={{ minHeight: 'calc(80vh - 200px)', maxHeight: 'calc(80vh - 150px)' }}
/>
<TokenCount>Tokens: {tokenCount}</TokenCount>
</TextAreaContainer>
<HStack width="100%" justifyContent="flex-end" mt="10px"> <HStack width="100%" justifyContent="flex-end" mt="10px">
<Button type="primary" onClick={onOk}> <Button type="primary" onClick={onOk}>
{t('common.close')} {t('common.close')}
@ -116,4 +129,21 @@ const EmojiButtonWrapper = styled.div`
} }
` `
const TextAreaContainer = styled.div`
position: relative;
width: 100%;
`
const TokenCount = styled.div`
position: absolute;
bottom: 8px;
right: 8px;
background-color: var(--color-background-soft);
padding: 2px 8px;
border-radius: 4px;
font-size: 12px;
color: var(--color-text-2);
user-select: none;
`
export default AssistantPromptSettings export default AssistantPromptSettings