chore: chat hooks
This commit is contained in:
jialin
@@ -0,0 +1,218 @@
|
||||
import useOverlayScroller from '@/hooks/use-overlay-scroller';
|
||||
import { fetchChunkedData, readStreamData } from '@/utils/fetch-chunk-data';
|
||||
import _ from 'lodash';
|
||||
import { useEffect, useRef, useState } from 'react';
|
||||
import { CHAT_API } from '../apis';
|
||||
import { Roles, generateMessages } from '../config';
|
||||
import { MessageItem } from '../config/types';
|
||||
|
||||
export default function useChatCompletion(
|
||||
scroller: React.RefObject<HTMLElement>
|
||||
) {
|
||||
const { initialize, updateScrollerPosition } = useOverlayScroller();
|
||||
const [loading, setLoading] = useState(false);
|
||||
const [tokenResult, setTokenResult] = useState<any>(null);
|
||||
const [messageList, setMessageList] = useState<MessageItem[]>([]);
|
||||
const controllerRef = useRef<any>(null);
|
||||
const messageId = useRef<number>(0);
|
||||
const contentRef = useRef<any>('');
|
||||
const currentMessageRef = useRef<any>(null);
|
||||
const messageListLengthCache = useRef<number>(0);
|
||||
const reasonContentRef = useRef('');
|
||||
|
||||
const setMessageId = () => {
|
||||
messageId.current = messageId.current + 1;
|
||||
return messageId.current;
|
||||
};
|
||||
|
||||
const formatContent = (data: {
|
||||
content: string;
|
||||
reasoningContent: string;
|
||||
}) => {
|
||||
if (data.reasoningContent && !data.content) {
|
||||
return `<think>${data.reasoningContent}`;
|
||||
}
|
||||
if (data.reasoningContent && data.content) {
|
||||
return `<think>${data.reasoningContent}</think>${data.content}`;
|
||||
}
|
||||
return data.content;
|
||||
};
|
||||
|
||||
const joinMessage = (chunk: any) => {
|
||||
console.log('chunk:', chunk);
|
||||
setTokenResult({
|
||||
...(chunk?.usage ?? {})
|
||||
});
|
||||
|
||||
if (!chunk || !_.get(chunk, 'choices', []).length) {
|
||||
return;
|
||||
}
|
||||
|
||||
reasonContentRef.current =
|
||||
reasonContentRef.current +
|
||||
_.get(chunk, 'choices.0.delta.reasoning_content', '');
|
||||
contentRef.current =
|
||||
contentRef.current + _.get(chunk, 'choices.0.delta.content', '');
|
||||
|
||||
const content = formatContent({
|
||||
content: contentRef.current,
|
||||
reasoningContent: reasonContentRef.current
|
||||
});
|
||||
|
||||
setMessageList([
|
||||
...messageList,
|
||||
...currentMessageRef.current,
|
||||
{
|
||||
role: Roles.Assistant,
|
||||
content: content,
|
||||
uid: messageId.current
|
||||
}
|
||||
]);
|
||||
};
|
||||
|
||||
const handleClear = () => {
|
||||
setMessageList([]);
|
||||
setTokenResult(null);
|
||||
};
|
||||
|
||||
const handleAddNewMessage = (message?: { role: string; content: string }) => {
|
||||
const newMessage = message || {
|
||||
role:
|
||||
_.last(messageList)?.role === Roles.User ? Roles.Assistant : Roles.User,
|
||||
content: ''
|
||||
};
|
||||
setMessageList((preList) => [
|
||||
...preList,
|
||||
{
|
||||
...newMessage,
|
||||
uid: setMessageId()
|
||||
}
|
||||
]);
|
||||
};
|
||||
|
||||
const handleStopConversation = () => {
|
||||
controllerRef.current?.abort?.();
|
||||
setLoading(false);
|
||||
};
|
||||
|
||||
const resetPerRequestCache = () => {
|
||||
contentRef.current = '';
|
||||
reasonContentRef.current = '';
|
||||
};
|
||||
|
||||
const submitMessage = async (params: {
|
||||
current?: { role: string; content: string };
|
||||
system?: { role: string; content: string };
|
||||
parameters: any;
|
||||
}) => {
|
||||
try {
|
||||
setLoading(true);
|
||||
setMessageId();
|
||||
setTokenResult(null);
|
||||
|
||||
const { current, parameters, system } = params;
|
||||
|
||||
controllerRef.current?.abort?.();
|
||||
controllerRef.current = new AbortController();
|
||||
const signal = controllerRef.current.signal;
|
||||
currentMessageRef.current = current
|
||||
? [
|
||||
{
|
||||
...current,
|
||||
uid: messageId.current
|
||||
}
|
||||
]
|
||||
: [];
|
||||
|
||||
resetPerRequestCache();
|
||||
setMessageList((pre) => {
|
||||
return [...pre, ...currentMessageRef.current];
|
||||
});
|
||||
|
||||
const messageParams = [
|
||||
...(system ? [system] : []),
|
||||
...messageList,
|
||||
...currentMessageRef.current
|
||||
];
|
||||
|
||||
const messages = generateMessages(messageParams);
|
||||
|
||||
const chatParams = {
|
||||
messages: messages,
|
||||
...parameters,
|
||||
stream: true,
|
||||
stream_options: {
|
||||
include_usage: true
|
||||
}
|
||||
};
|
||||
const result: any = await fetchChunkedData({
|
||||
data: chatParams,
|
||||
url: CHAT_API,
|
||||
signal
|
||||
});
|
||||
if (result?.error) {
|
||||
setTokenResult({
|
||||
error: true,
|
||||
errorMessage:
|
||||
result?.data?.error?.message || result?.data?.message || ''
|
||||
});
|
||||
return;
|
||||
}
|
||||
setMessageId();
|
||||
const { reader, decoder } = result;
|
||||
await readStreamData(reader, decoder, (chunk: any) => {
|
||||
if (chunk?.error) {
|
||||
setTokenResult({
|
||||
error: true,
|
||||
errorMessage: chunk?.error?.message || chunk?.message || ''
|
||||
});
|
||||
return;
|
||||
}
|
||||
joinMessage(chunk);
|
||||
});
|
||||
} catch (error) {
|
||||
console.log('error:', error);
|
||||
} finally {
|
||||
setLoading(false);
|
||||
}
|
||||
};
|
||||
const throttleUpdatePosition = _.throttle(updateScrollerPosition, 100);
|
||||
|
||||
useEffect(() => {
|
||||
if (scroller.current) {
|
||||
initialize(scroller.current);
|
||||
}
|
||||
}, [scroller.current, initialize]);
|
||||
|
||||
useEffect(() => {
|
||||
if (loading) {
|
||||
updateScrollerPosition();
|
||||
}
|
||||
}, [messageList, loading]);
|
||||
|
||||
useEffect(() => {
|
||||
if (messageList.length > messageListLengthCache.current) {
|
||||
updateScrollerPosition();
|
||||
}
|
||||
messageListLengthCache.current = messageList.length;
|
||||
}, [messageList.length]);
|
||||
|
||||
useEffect(() => {
|
||||
return () => {
|
||||
handleStopConversation();
|
||||
};
|
||||
}, []);
|
||||
|
||||
return {
|
||||
loading,
|
||||
tokenResult,
|
||||
messageList,
|
||||
setMessageId,
|
||||
setMessageList,
|
||||
handleClear,
|
||||
handleAddNewMessage,
|
||||
handleStopConversation,
|
||||
updateScrollerPosition,
|
||||
submitMessage
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1,117 @@
|
||||
import _ from 'lodash';
|
||||
import { imageSizeOptions } from '../config/params-config';
|
||||
|
||||
const LLM_METAKEYS: Record<string, any> = {
|
||||
seed: 'seed',
|
||||
stop: 'stop',
|
||||
temperature: 'temperature',
|
||||
top_p: 'top_p',
|
||||
n_ctx: 'n_ctx',
|
||||
n_slot: 'n_slot',
|
||||
max_model_len: 'max_model_len'
|
||||
};
|
||||
|
||||
const IMG_METAKEYS = [
|
||||
'sample_method',
|
||||
'sampling_steps',
|
||||
'schedule_method',
|
||||
'cfg_scale',
|
||||
'guidance',
|
||||
'negative_prompt'
|
||||
];
|
||||
|
||||
const llmInitialValues = {
|
||||
seed: null,
|
||||
stop: null,
|
||||
temperature: 1,
|
||||
top_p: 1,
|
||||
max_tokens: 1024
|
||||
};
|
||||
|
||||
const imgInitialValues = {
|
||||
n: 1,
|
||||
seed: null,
|
||||
sample_method: 'euler_a',
|
||||
cfg_scale: 4.5,
|
||||
guidance: 3.5,
|
||||
sampling_steps: 10,
|
||||
negative_prompt: null,
|
||||
schedule_method: 'discrete',
|
||||
preview: 'preview_faster'
|
||||
};
|
||||
|
||||
export default function useInitMeta() {
|
||||
const extractLLMMeta = (meta: any) => {
|
||||
const modelMeta = meta || {};
|
||||
const modelMetaValue = _.pick(modelMeta, _.keys(LLM_METAKEYS));
|
||||
const obj = Object.entries(LLM_METAKEYS).reduce(
|
||||
(acc: any, [key, value]) => {
|
||||
const val = modelMetaValue[key];
|
||||
if (val && _.hasIn(modelMetaValue, key)) {
|
||||
acc[value] = val;
|
||||
}
|
||||
return acc;
|
||||
},
|
||||
{}
|
||||
);
|
||||
|
||||
let defaultMaxTokens = 1024;
|
||||
|
||||
if (obj.n_ctx && obj.n_slot) {
|
||||
defaultMaxTokens = _.divide(obj.n_ctx / 2, obj.n_slot);
|
||||
} else if (obj.max_model_len) {
|
||||
defaultMaxTokens = obj.max_model_len / 2;
|
||||
}
|
||||
|
||||
return {
|
||||
form: _.merge({}, llmInitialValues, {
|
||||
..._.omit(obj, ['n_ctx', 'n_slot', 'max_model_len']),
|
||||
max_tokens: defaultMaxTokens
|
||||
}),
|
||||
meta: {
|
||||
...obj,
|
||||
max_tokens: obj.max_model_len || _.divide(obj.n_ctx, obj.n_slot)
|
||||
}
|
||||
};
|
||||
};
|
||||
|
||||
const getNewImageSizeOptions = (metaData: any) => {
|
||||
const { max_height, max_width } = metaData || {};
|
||||
if (!max_height || !max_width) {
|
||||
return imageSizeOptions;
|
||||
}
|
||||
const newImageSizeOptions = imageSizeOptions.filter((item) => {
|
||||
return item.width <= max_width && item.height <= max_height;
|
||||
});
|
||||
if (
|
||||
!newImageSizeOptions.find(
|
||||
(item) => item.width === max_width && item.height === max_height
|
||||
)
|
||||
) {
|
||||
newImageSizeOptions.push({
|
||||
width: max_width,
|
||||
height: max_height,
|
||||
label: `${max_width}x${max_height}`,
|
||||
value: `${max_width}x${max_height}`
|
||||
});
|
||||
}
|
||||
return newImageSizeOptions;
|
||||
};
|
||||
|
||||
const extractIMGMeta = (meta: any) => {
|
||||
return {
|
||||
form: _.merge({}, imgInitialValues, {
|
||||
..._.pick(meta, IMG_METAKEYS),
|
||||
width: meta?.default_width || 512,
|
||||
height: meta?.default_height || 512
|
||||
}),
|
||||
meta: meta,
|
||||
sizeOptions: getNewImageSizeOptions(meta)
|
||||
};
|
||||
};
|
||||
|
||||
return {
|
||||
extractLLMMeta,
|
||||
extractIMGMeta
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1,199 @@
|
||||
import { CREAT_IMAGE_API } from '@/pages/playground/apis';
|
||||
import { extractErrorMessage, promptList } from '@/pages/playground/config';
|
||||
import { generateRandomNumber } from '@/utils';
|
||||
import {
|
||||
fetchChunkedData,
|
||||
readLargeStreamData as readStreamData
|
||||
} from '@/utils/fetch-chunk-data';
|
||||
import _ from 'lodash';
|
||||
import { useCallback, useRef, useState } from 'react';
|
||||
|
||||
const ODD_STRING = 'AAAABJRU5ErkJgg===';
|
||||
|
||||
export default function useTextImage() {
|
||||
const [loading, setLoading] = useState(false);
|
||||
const [tokenResult, setTokenResult] = useState<any>(null);
|
||||
const [imageList, setImageList] = useState<any[]>([]);
|
||||
const [currentPrompt, setCurrentPrompt] = useState('');
|
||||
const messageId = useRef<number>(0);
|
||||
const requestToken = useRef<any>(null);
|
||||
|
||||
const removeBase64Suffix = (str: string, suffix: string) => {
|
||||
return str.endsWith(suffix) ? str.slice(0, -suffix.length) : str;
|
||||
};
|
||||
|
||||
const setImageSize = useCallback((parameters: any) => {
|
||||
let size: Record<string, string | number> = {
|
||||
span: 12
|
||||
};
|
||||
if (parameters.n === 1) {
|
||||
size.span = 24;
|
||||
}
|
||||
if (parameters.n === 2) {
|
||||
size.span = 12;
|
||||
}
|
||||
if (parameters.n === 3) {
|
||||
size.span = 12;
|
||||
}
|
||||
if (parameters.n === 4) {
|
||||
size.span = 12;
|
||||
}
|
||||
return size;
|
||||
}, []);
|
||||
|
||||
const setMessageId = () => {
|
||||
messageId.current = messageId.current + 1;
|
||||
return messageId.current;
|
||||
};
|
||||
|
||||
const generateNumber = (min: number, max: number) => {
|
||||
return Math.floor(Math.random() * (max - min + 1) + min);
|
||||
};
|
||||
|
||||
const submitMessage = async (params: {
|
||||
current?: { content: string };
|
||||
system?: { role: string; content: string };
|
||||
parameters: any;
|
||||
}) => {
|
||||
const { current, parameters } = params;
|
||||
try {
|
||||
if (!parameters.model) return;
|
||||
const size: any = setImageSize(parameters);
|
||||
setLoading(true);
|
||||
setMessageId();
|
||||
setTokenResult(null);
|
||||
setCurrentPrompt(current?.content || '');
|
||||
|
||||
const imgSize = [parameters.width, parameters.height];
|
||||
|
||||
// preview
|
||||
let stream_options: Record<string, any> = {
|
||||
chunk_size: 16 * 1024,
|
||||
chunk_results: true
|
||||
};
|
||||
if (parameters.preview === 'preview') {
|
||||
stream_options = {
|
||||
preview: true
|
||||
};
|
||||
}
|
||||
|
||||
if (parameters.preview === 'preview_faster') {
|
||||
stream_options = {
|
||||
preview_faster: true
|
||||
};
|
||||
}
|
||||
|
||||
let newImageList = Array(parameters.n)
|
||||
.fill({})
|
||||
.map((item, index: number) => {
|
||||
return {
|
||||
dataUrl: 'data:image/png;base64,',
|
||||
...size,
|
||||
progress: 0,
|
||||
height: imgSize[1],
|
||||
width: imgSize[0],
|
||||
loading: true,
|
||||
progressType: 'dashboard',
|
||||
preview: false,
|
||||
uid: setMessageId()
|
||||
};
|
||||
});
|
||||
setImageList(newImageList);
|
||||
|
||||
requestToken.current?.abort?.();
|
||||
requestToken.current = new AbortController();
|
||||
|
||||
const params = {
|
||||
..._.omitBy(
|
||||
parameters,
|
||||
(value: string, key: string) =>
|
||||
!value || ['width', 'height', 'seed'].includes(key)
|
||||
),
|
||||
size: `${imgSize[0]}x${imgSize[1]}`,
|
||||
seed: parameters.random_seed ? generateRandomNumber() : parameters.seed,
|
||||
stream: true,
|
||||
stream_options: {
|
||||
...stream_options
|
||||
},
|
||||
prompt: current?.content
|
||||
};
|
||||
|
||||
const result: any = await fetchChunkedData({
|
||||
data: params,
|
||||
url: `${CREAT_IMAGE_API}?t=${Date.now()}`,
|
||||
signal: requestToken.current.signal
|
||||
});
|
||||
if (result.error) {
|
||||
setTokenResult({
|
||||
error: true,
|
||||
errorMessage: extractErrorMessage(result)
|
||||
});
|
||||
setImageList([]);
|
||||
return;
|
||||
}
|
||||
|
||||
const { reader, decoder } = result;
|
||||
|
||||
await readStreamData(reader, decoder, (chunk: any) => {
|
||||
if (chunk?.error) {
|
||||
setTokenResult({
|
||||
error: true,
|
||||
errorMessage: chunk?.error?.message || chunk?.message || ''
|
||||
});
|
||||
return;
|
||||
}
|
||||
chunk?.data?.forEach((item: any) => {
|
||||
const imgItem = newImageList[item.index];
|
||||
if (item.b64_json && stream_options.chunk_results) {
|
||||
imgItem.dataUrl += removeBase64Suffix(item.b64_json, ODD_STRING);
|
||||
} else if (item.b64_json) {
|
||||
imgItem.dataUrl = `data:image/png;base64,${removeBase64Suffix(item.b64_json, ODD_STRING)}`;
|
||||
}
|
||||
const progress = item.progress;
|
||||
|
||||
newImageList[item.index] = {
|
||||
dataUrl: imgItem.dataUrl,
|
||||
height: imgSize[1],
|
||||
width: imgSize[0],
|
||||
maxHeight: `${imgSize[1]}px`,
|
||||
maxWidth: `${imgSize[0]}px`,
|
||||
uid: imgItem.uid,
|
||||
span: imgItem.span,
|
||||
loading: stream_options.chunk_results ? progress < 100 : false,
|
||||
preview: progress >= 100,
|
||||
progress: progress
|
||||
};
|
||||
});
|
||||
setImageList([...newImageList]);
|
||||
});
|
||||
} catch (error) {
|
||||
console.log('error:', error);
|
||||
requestToken.current?.abort?.();
|
||||
setImageList([]);
|
||||
} finally {
|
||||
setLoading(false);
|
||||
}
|
||||
};
|
||||
|
||||
const handleClear = () => {
|
||||
setMessageId();
|
||||
setImageList([]);
|
||||
setTokenResult(null);
|
||||
};
|
||||
|
||||
const handleStopConversation = () => {
|
||||
requestToken.current?.abort?.();
|
||||
setLoading(false);
|
||||
};
|
||||
|
||||
return {
|
||||
loading,
|
||||
tokenResult,
|
||||
imageList,
|
||||
promptList,
|
||||
handleStopConversation,
|
||||
generateNumber,
|
||||
handleClear,
|
||||
submitMessage
|
||||
};
|
||||
}
|
||||
Reference in New Issue
Block a user