import useSelectRender from '@/components/seal-form/hooks/use-select-render'; import { queryModelsList } from '@/pages/llmodels/apis'; import { ListItem as ModelListItem } from '@/pages/llmodels/config/types'; import { queryUsersList } from '@/pages/users/apis'; import { DownloadOutlined } from '@ant-design/icons'; import { useIntl } from '@umijs/max'; import { Button, DatePicker, Select, Tooltip } from 'antd'; import dayjs from 'dayjs'; import _ from 'lodash'; import { useMemo, useState } from 'react'; import styled from 'styled-components'; import { DASHBOARD_STATS_API, DASHBOARD_USAGE_API, queryDashboardUsageData } from '../../apis'; import { baseColorMap } from '../../config'; import useRangePickerPreset from '../../hooks/use-rangepicker-preset'; const FilterWrapper = styled.div` display: flex; justify-content: space-between; align-items: center; margin-bottom: 0px; .selection { display: flex; align-items: center; gap: 12px; } `; interface RequestTokenData { requestData: { name: string; color: string; areaStyle: any; data: { time: string; value: number }[]; }[]; tokenData: { data: { time: string; value: number }[]; }[]; xAxisData: string[]; } interface TopUserData { userData: { name: string; value: number }[]; topUserList: string[]; } const getLast30Days = () => { const dates: string[] = []; for (let i = 29; i >= 0; i--) { const date = dayjs().subtract(i, 'day').format('YYYY-MM-DD'); dates.push(date); } return dates; }; const getAllDays = (start: string, end: string) => { const startDate = dayjs(start); const endDate = dayjs(end); const days: string[] = []; for ( let d = startDate; d.isBefore(endDate) || d.isSame(endDate, 'day'); d = d.add(1, 'day') ) { days.push(d.format('YYYY-MM-DD')); } return days; }; const generateValueMap = (list: { timestamp: number; value: number }[]) => { return new Map( list.map((item) => [ dayjs(item.timestamp * 1000).format('YYYY-MM-DD'), item.value ]) ); }; const generateData = (dateRage: string[], valueMap: Map) => { return dateRage.map((date) => { const value = valueMap.get(date) || 0; return { time: date, value: value }; }); }; export default function useUseageData(config: { url: string; defaultData?: T; }) { const { url, defaultData } = config || {}; const intl = useIntl(); const { TagRender } = useSelectRender({ maxTagWidth: 100 }); const { disabledRangeDaysDate, rangePresets } = useRangePickerPreset({ range: 60 }); const [open, setOpen] = useState(false); const [result, setResult] = useState<{ start_date?: string; end_date?: string; data: T; }>({ start_date: dayjs().subtract(30, 'days').format('YYYY-MM-DD'), end_date: dayjs().format('YYYY-MM-DD'), data: defaultData as T }); const [query, setQuery] = useState<{ start_date: string; end_date: string; model_ids: number[]; user_ids: number[]; }>({ start_date: dayjs().subtract(30, 'days').format('YYYY-MM-DD'), end_date: dayjs().format('YYYY-MM-DD'), model_ids: [], user_ids: [] }); const [modelList, setModelList] = useState[]>([]); const [userList, setUserList] = useState[]>([]); const [loading, setLoading] = useState(false); const usageData = useMemo<{ requestTokenData: RequestTokenData; }>(() => { if (url === DASHBOARD_USAGE_API) { return { requestTokenData: { requestData: [], tokenData: [], xAxisData: [] } }; } const { start_date, end_date, data } = result as { start_date?: string; end_date?: string; data: { api_request_history: { timestamp: number; value: number }[]; completion_token_history: { timestamp: number; value: number }[]; prompt_token_history: { timestamp: number; value: number }[]; }; }; const dateRange = start_date && end_date ? getAllDays(start_date, end_date) : getLast30Days(); const completionTokenHistory = data?.completion_token_history || []; const promptTokenHistory = data?.prompt_token_history || []; const apiRequestHistory = data?.api_request_history || []; if (!completionTokenHistory.length) { return { requestTokenData: { requestData: [], tokenData: [], xAxisData: [] } }; } // ========== API request ============== const requestList: { name: string; color: string; areaStyle: any; data: { time: string; value: number }[]; } = { name: 'API requests', areaStyle: { color: 'rgba(13,171,219,0.15)' }, color: baseColorMap.baseR1, data: generateData(dateRange, generateValueMap(apiRequestHistory)) }; // =========== token usage data ============== const completionData: any = { name: 'Completion tokens', color: baseColorMap.base, data: generateData(dateRange, generateValueMap(completionTokenHistory)) }; const promptData: any = { name: 'Prompt tokens', color: baseColorMap.baseR3, data: generateData(dateRange, generateValueMap(promptTokenHistory)) }; return { requestTokenData: { requestData: [requestList], tokenData: [completionData, promptData], xAxisData: dateRange } }; }, [result, url]); const fetchModelsList = async () => { try { const params = { page: 1, page_size: 100 }; const response = await queryModelsList(params); const list = _.map(response.items || [], (item: ModelListItem) => { return { label: item.name, value: item.id }; }); setModelList(list); } catch (error) { setModelList([]); } }; const fetchUsersList = async () => { try { const params = { page: 1, page_size: 100 }; const response = await queryUsersList(params); const list = _.map(response.items || [], (item: any) => { return { label: item.username, value: item.id }; }); setUserList(list); } catch (error) { setUserList([]); } }; const fetchUsageData = async (queryParams: any) => { try { setLoading(true); const response = await queryDashboardUsageData(queryParams, { url: url }); setResult({ start_date: queryParams.start_date, end_date: queryParams.end_date, data: response }); } catch (error) { setResult({ start_date: queryParams.start_date, end_date: queryParams.end_date, data: {} as T }); } finally { setLoading(false); } }; const handleDateChange = (dates: any, dateString: string[]) => { setQuery((pre) => { return { ...pre, start_date: dateString[0], end_date: dateString[1] }; }); fetchUsageData({ ...query, start_date: dateString[0], end_date: dateString[1] }); }; const handleOnCancel = () => { setOpen(false); }; const handleExport = () => { setOpen(true); }; const handleUsersChange = (value: number[]) => { setQuery((pre) => { return { ...pre, user_ids: value }; }); fetchUsageData({ ...query, user_ids: value }); }; const handleModelsChange = (value: number[]) => { setQuery((pre) => { return { ...pre, model_ids: value }; }); fetchUsageData({ ...query, model_ids: value }); }; const init = () => { fetchUsageData(query); fetchModelsList(); fetchUsersList(); }; const FilterBar = () => { return (
{url === DASHBOARD_STATS_API && ( )}
); }; return { usageData, result, open, loading, userList, modelList, query, setQuery, init, setResult, FilterBar, handleOnCancel, handleExport }; }