Spaces:
Running
Running
| import React, { createContext, useState, useContext, ReactNode, useEffect } from 'react'; | |
| type DeploymentGroupContextProps = { | |
| deploymentGroups: string[]; | |
| fetchDeploymentGroups: () => void; | |
| refreshing: boolean; | |
| ready: boolean; | |
| }; | |
| type DeploymentProps = { | |
| children: ReactNode; | |
| }; | |
| type Reader = { | |
| consumer_group: string; | |
| model_type: string; | |
| }; | |
| interface Data { | |
| live_readers: { [key: string]: Reader }; | |
| } | |
| const DeploymentGroupContext = createContext<DeploymentGroupContextProps | undefined>(undefined); | |
| export const DeploymentGroupProvider = ({ children }: DeploymentProps) => { | |
| const [deploymentGroups, setDeploymentGroups] = useState<string[]>([]); | |
| const [refreshing, setRefreshing] = useState<boolean>(true); | |
| const [ready, setReady] = useState<boolean>(false); | |
| const isStaticDeployment = process.env.NODE_ENV === 'production'; // Adjust this based on your environment detection | |
| async function fetchDeploymentGroups() { | |
| if (isStaticDeployment) { | |
| // Skip the fetch and status checking in a static environment | |
| setRefreshing(false); | |
| setDeploymentGroups([]); | |
| setReady(true); | |
| return; | |
| } | |
| const url = 'status'; | |
| setRefreshing(true); | |
| try { | |
| const response = await fetch(url, { | |
| method: 'GET', | |
| headers: { | |
| 'Content-Type': 'application/json', | |
| }, | |
| }); | |
| const data: Data = await response.json(); | |
| // Logic for filtering consumer groups | |
| const groupMap = new Map(); | |
| const groups = Object.values(data['live_readers']).reduce( | |
| (accumulator: string[], reader: Reader) => { | |
| if (reader.model_type === 'EMBEDDING') { | |
| groupMap.set(reader.consumer_group, true); | |
| } else if (!groupMap.has(reader.consumer_group)) { | |
| accumulator.push(reader.consumer_group); | |
| } | |
| return accumulator; | |
| }, | |
| [] | |
| ); | |
| groupMap.forEach((value, key) => { | |
| if (value) { | |
| const index = groups.indexOf(key); | |
| if (index !== -1) { | |
| groups.splice(index, 1); | |
| } | |
| } | |
| }); | |
| setDeploymentGroups(groups); | |
| if (groups.length > 0) { | |
| setReady(true); | |
| } | |
| setRefreshing(false); | |
| } catch (err) { | |
| setRefreshing(false); | |
| console.log(err); | |
| } | |
| } | |
| useEffect(() => { | |
| if (isStaticDeployment) { | |
| // Skip polling in a static environment | |
| return; | |
| } | |
| const pollingInterval = setInterval(() => { | |
| if (!ready) { | |
| fetchDeploymentGroups(); | |
| } | |
| }, 1000); | |
| if (ready) { | |
| clearInterval(pollingInterval); | |
| } | |
| return () => clearInterval(pollingInterval); | |
| }, [ready, isStaticDeployment]); | |
| return ( | |
| <DeploymentGroupContext.Provider | |
| value={{ deploymentGroups, fetchDeploymentGroups, refreshing, ready }} | |
| > | |
| {children} | |
| </DeploymentGroupContext.Provider> | |
| ); | |
| }; | |
| export const useDeploymentGroup = () => { | |
| const context = useContext(DeploymentGroupContext); | |
| if (!context) { | |
| throw new Error('useDeploymentGroup must be used within a DeploymentGroupProvider'); | |
| } | |
| return context; | |
| }; | |