refactor: useMergedWorkers

feat-display-remote-workers
liangfung 2023-11-21 23:10:46 +08:00
parent 0038d8598e
commit c952819067
1 changed files with 21 additions and 13 deletions

View File

@ -1,4 +1,4 @@
import { groupBy, findIndex } from 'lodash-es' import { groupBy, findIndex, slice } from 'lodash-es'
import { import {
GetWorkersDocument, GetWorkersDocument,
Worker, Worker,
@ -6,6 +6,7 @@ import {
} from '@/lib/gql/generates/graphql' } from '@/lib/gql/generates/graphql'
import { useGraphQL } from './use-graphql' import { useGraphQL } from './use-graphql'
import type { HealthInfo } from './use-health' import type { HealthInfo } from './use-health'
import React from 'react'
function useRemoteWorkers() { function useRemoteWorkers() {
return useGraphQL(GetWorkersDocument) return useGraphQL(GetWorkersDocument)
@ -33,20 +34,27 @@ function transformHealthInfoToWorker(
function useMergedWorkers(healthInfo: HealthInfo | undefined) { function useMergedWorkers(healthInfo: HealthInfo | undefined) {
const { data } = useRemoteWorkers() const { data } = useRemoteWorkers()
let workers = data?.workers || [] let workers = data?.workers
const groupedWorkers = React.useMemo(() => {
const _workers = slice(workers)
const haveRemoteCompletionWorkers = const haveRemoteCompletionWorkers =
findIndex(workers, { kind: WorkerKind.Completion }) > -1 findIndex(workers, { kind: WorkerKind.Completion }) > -1
const haveRemoteChatWorkers = const haveRemoteChatWorkers =
findIndex(workers, { kind: WorkerKind.Chat }) > -1 findIndex(workers, { kind: WorkerKind.Chat }) > -1
if (!haveRemoteCompletionWorkers && healthInfo?.model) { if (!haveRemoteCompletionWorkers && healthInfo?.model) {
workers.push(transformHealthInfoToWorker(healthInfo, WorkerKind.Completion)) _workers.push(
transformHealthInfoToWorker(healthInfo, WorkerKind.Completion)
)
} }
if (!haveRemoteChatWorkers && healthInfo?.chat_model) { if (!haveRemoteChatWorkers && healthInfo?.chat_model) {
workers.push(transformHealthInfoToWorker(healthInfo, WorkerKind.Chat)) _workers.push(transformHealthInfoToWorker(healthInfo, WorkerKind.Chat))
} }
return groupBy(workers, worker => worker.kind) return groupBy(_workers, 'kind')
}, [healthInfo, workers])
return groupedWorkers
} }
export { useRemoteWorkers, useMergedWorkers } export { useRemoteWorkers, useMergedWorkers }