feat: preserve token speed in the thread (#4711)

* feat: preserve token speed in the thread

* chore: lint fix
This commit is contained in:
Louis 2025-02-21 00:45:11 +07:00 committed by GitHub
parent fddb7251fb
commit cddaf61c99
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 27 additions and 8 deletions

View File

@ -18,7 +18,7 @@ import {
extractInferenceParams, extractInferenceParams,
ModelExtension, ModelExtension,
} from '@janhq/core' } from '@janhq/core'
import { useAtomValue, useSetAtom } from 'jotai' import { useAtom, useAtomValue, useSetAtom } from 'jotai'
import { ulid } from 'ulidx' import { ulid } from 'ulidx'
import { activeModelAtom, stateModelAtom } from '@/hooks/useActiveModel' import { activeModelAtom, stateModelAtom } from '@/hooks/useActiveModel'
@ -75,8 +75,10 @@ export default function ModelHandler() {
const activeThreadRef = useRef(activeThread) const activeThreadRef = useRef(activeThread)
const activeModelParams = useAtomValue(getActiveThreadModelParamsAtom) const activeModelParams = useAtomValue(getActiveThreadModelParamsAtom)
const activeModelParamsRef = useRef(activeModelParams) const activeModelParamsRef = useRef(activeModelParams)
const setTokenSpeed = useSetAtom(tokenSpeedAtom)
const [tokenSpeed, setTokenSpeed] = useAtom(tokenSpeedAtom)
const { engines } = useGetEngines() const { engines } = useGetEngines()
const tokenSpeedRef = useRef(tokenSpeed)
useEffect(() => { useEffect(() => {
activeThreadRef.current = activeThread activeThreadRef.current = activeThread
@ -106,6 +108,10 @@ export default function ModelHandler() {
messageGenerationSubscriber.current = subscribedGeneratingMessage messageGenerationSubscriber.current = subscribedGeneratingMessage
}, [subscribedGeneratingMessage]) }, [subscribedGeneratingMessage])
useEffect(() => {
tokenSpeedRef.current = tokenSpeed
}, [tokenSpeed])
const onNewMessageResponse = useCallback( const onNewMessageResponse = useCallback(
async (message: ThreadMessage) => { async (message: ThreadMessage) => {
if (message.type === MessageRequestType.Thread) { if (message.type === MessageRequestType.Thread) {
@ -275,6 +281,12 @@ export default function ModelHandler() {
metadata, metadata,
}) })
// Update message's metadata with token usage
message.metadata = {
...message.metadata,
token_speed: tokenSpeedRef.current?.tokenSpeed,
}
if (message.status === MessageStatus.Error) { if (message.status === MessageStatus.Error) {
message.metadata = { message.metadata = {
...message.metadata, ...message.metadata,

View File

@ -122,13 +122,20 @@ const MessageContainer: React.FC<
)} )}
> >
<div> <div>
{tokenSpeed && {((!!tokenSpeed &&
tokenSpeed.message === props.id && tokenSpeed.message === props.id &&
tokenSpeed.tokenSpeed > 0 && ( tokenSpeed.tokenSpeed > 0) ||
<p className="text-xs font-medium text-[hsla(var(--text-secondary))]"> (props.metadata &&
Token Speed: {Number(tokenSpeed.tokenSpeed).toFixed(2)}t/s 'token_speed' in props.metadata &&
</p> !!props.metadata?.token_speed)) && (
)} <p className="text-xs font-medium text-[hsla(var(--text-secondary))]">
Token Speed:{' '}
{Number(
props.metadata?.token_speed ?? tokenSpeed?.tokenSpeed
).toFixed(2)}
t/s
</p>
)}
</div> </div>
<MessageToolbar message={props} /> <MessageToolbar message={props} />