diff options
| author | Taylor Mullen <[email protected]> | 2025-04-17 18:06:21 -0400 |
|---|---|---|
| committer | N. Taylor Mullen <[email protected]> | 2025-04-17 15:29:34 -0700 |
| commit | cfc697a96d2e716a75e1c3b7f0f34fce81abaf1e (patch) | |
| tree | e06bcba67ca71a874048aa887b17457dbd409bdf /packages/cli/src/ui/hooks | |
| parent | 7928c1727f0b208ed34850cc89bbb36ea3dd23e5 (diff) | |
Run `npm run format`
- Also updated README.md accordingly.
Part of https://b.corp.google.com/issues/411384603
Diffstat (limited to 'packages/cli/src/ui/hooks')
| -rw-r--r-- | packages/cli/src/ui/hooks/useGeminiStream.ts | 243 | ||||
| -rw-r--r-- | packages/cli/src/ui/hooks/useLoadingIndicator.ts | 100 |
2 files changed, 186 insertions, 157 deletions
diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index c65422ca..4144d96a 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -7,136 +7,157 @@ import { processGeminiStream } from '../../core/gemini-stream.js'; import { StreamingState } from '../../core/gemini-stream.js'; const addHistoryItem = ( - setHistory: React.Dispatch<React.SetStateAction<HistoryItem[]>>, - itemData: Omit<HistoryItem, 'id'>, - id: number + setHistory: React.Dispatch<React.SetStateAction<HistoryItem[]>>, + itemData: Omit<HistoryItem, 'id'>, + id: number, ) => { - setHistory((prevHistory) => [ - ...prevHistory, - { ...itemData, id } as HistoryItem, - ]); + setHistory((prevHistory) => [ + ...prevHistory, + { ...itemData, id } as HistoryItem, + ]); }; export const useGeminiStream = ( - setHistory: React.Dispatch<React.SetStateAction<HistoryItem[]>>, + setHistory: React.Dispatch<React.SetStateAction<HistoryItem[]>>, ) => { - const [streamingState, setStreamingState] = useState<StreamingState>(StreamingState.Idle); - const [initError, setInitError] = useState<string | null>(null); - const abortControllerRef = useRef<AbortController | null>(null); - const currentToolGroupIdRef = useRef<number | null>(null); - const chatSessionRef = useRef<Chat | null>(null); - const geminiClientRef = useRef<GeminiClient | null>(null); - const messageIdCounterRef = useRef(0); + const [streamingState, setStreamingState] = useState<StreamingState>( + StreamingState.Idle, + ); + const [initError, setInitError] = useState<string | null>(null); + const abortControllerRef = useRef<AbortController | null>(null); + const currentToolGroupIdRef = useRef<number | null>(null); + const chatSessionRef = useRef<Chat | null>(null); + const geminiClientRef = useRef<GeminiClient | null>(null); + const messageIdCounterRef = useRef(0); - // Initialize Client Effect (remains the same) - useEffect(() => { - setInitError(null); - if (!geminiClientRef.current) { - try { - geminiClientRef.current = new GeminiClient(); - } catch (error: any) { - setInitError(`Failed to initialize client: ${error.message || 'Unknown error'}`); - } - } - }, []); + // Initialize Client Effect (remains the same) + useEffect(() => { + setInitError(null); + if (!geminiClientRef.current) { + try { + geminiClientRef.current = new GeminiClient(); + } catch (error: any) { + setInitError( + `Failed to initialize client: ${error.message || 'Unknown error'}`, + ); + } + } + }, []); - // Input Handling Effect (remains the same) - useInput((input, key) => { - if (streamingState === StreamingState.Responding && key.escape) { - abortControllerRef.current?.abort(); - } - }); + // Input Handling Effect (remains the same) + useInput((input, key) => { + if (streamingState === StreamingState.Responding && key.escape) { + abortControllerRef.current?.abort(); + } + }); - // ID Generation Callback (remains the same) - const getNextMessageId = useCallback((baseTimestamp: number): number => { - messageIdCounterRef.current += 1; - return baseTimestamp + messageIdCounterRef.current; - }, []); + // ID Generation Callback (remains the same) + const getNextMessageId = useCallback((baseTimestamp: number): number => { + messageIdCounterRef.current += 1; + return baseTimestamp + messageIdCounterRef.current; + }, []); - // Submit Query Callback (updated to call processGeminiStream) - const submitQuery = useCallback(async (query: PartListUnion) => { - if (streamingState === StreamingState.Responding) { - // No-op if already going. - return; - } + // Submit Query Callback (updated to call processGeminiStream) + const submitQuery = useCallback( + async (query: PartListUnion) => { + if (streamingState === StreamingState.Responding) { + // No-op if already going. + return; + } - if (typeof query === 'string' && query.toString().trim().length === 0) { - return; - } + if (typeof query === 'string' && query.toString().trim().length === 0) { + return; + } - const userMessageTimestamp = Date.now(); - const client = geminiClientRef.current; - if (!client) { - setInitError("Gemini client is not available."); - return; - } + const userMessageTimestamp = Date.now(); + const client = geminiClientRef.current; + if (!client) { + setInitError('Gemini client is not available.'); + return; + } - if (!chatSessionRef.current) { - chatSessionRef.current = await client.startChat(); - } + if (!chatSessionRef.current) { + chatSessionRef.current = await client.startChat(); + } - // Reset state - setStreamingState(StreamingState.Responding); - setInitError(null); - currentToolGroupIdRef.current = null; - messageIdCounterRef.current = 0; - const chat = chatSessionRef.current; + // Reset state + setStreamingState(StreamingState.Responding); + setInitError(null); + currentToolGroupIdRef.current = null; + messageIdCounterRef.current = 0; + const chat = chatSessionRef.current; - try { - // Add user message - if (typeof query === 'string') { - const trimmedQuery = query.toString(); - addHistoryItem(setHistory, { type: 'user', text: trimmedQuery }, userMessageTimestamp); - } else if ( - // HACK to detect errored function responses. - typeof query === 'object' && - query !== null && - !Array.isArray(query) && // Ensure it's a single Part object - 'functionResponse' in query && // Check if it's a function response Part - query.functionResponse?.response && // Check if response object exists - 'error' in query.functionResponse.response // Check specifically for the 'error' key - ) { - const history = chat.getHistory(); - history.push({ role: 'user', parts: [query] }); - return; - } + try { + // Add user message + if (typeof query === 'string') { + const trimmedQuery = query.toString(); + addHistoryItem( + setHistory, + { type: 'user', text: trimmedQuery }, + userMessageTimestamp, + ); + } else if ( + // HACK to detect errored function responses. + typeof query === 'object' && + query !== null && + !Array.isArray(query) && // Ensure it's a single Part object + 'functionResponse' in query && // Check if it's a function response Part + query.functionResponse?.response && // Check if response object exists + 'error' in query.functionResponse.response // Check specifically for the 'error' key + ) { + const history = chat.getHistory(); + history.push({ role: 'user', parts: [query] }); + return; + } - // Prepare for streaming - abortControllerRef.current = new AbortController(); - const signal = abortControllerRef.current.signal; + // Prepare for streaming + abortControllerRef.current = new AbortController(); + const signal = abortControllerRef.current.signal; - // --- Delegate to Stream Processor --- + // --- Delegate to Stream Processor --- - const stream = client.sendMessageStream(chat, query, signal); + const stream = client.sendMessageStream(chat, query, signal); - const addHistoryItemFromStream = (itemData: Omit<HistoryItem, 'id'>, id: number) => { - addHistoryItem(setHistory, itemData, id); - }; - const getStreamMessageId = () => getNextMessageId(userMessageTimestamp); + const addHistoryItemFromStream = ( + itemData: Omit<HistoryItem, 'id'>, + id: number, + ) => { + addHistoryItem(setHistory, itemData, id); + }; + const getStreamMessageId = () => getNextMessageId(userMessageTimestamp); - // Call the renamed processor function - await processGeminiStream({ - stream, - signal, - setHistory, - submitQuery, - getNextMessageId: getStreamMessageId, - addHistoryItem: addHistoryItemFromStream, - currentToolGroupIdRef, - }); - } catch (error: any) { - // (Error handling for stream initiation remains the same) - console.error("Error initiating stream:", error); - if (error.name !== 'AbortError') { - // Use historyUpdater's function potentially? Or keep addHistoryItem here? - // Keeping addHistoryItem here for direct errors from this scope. - addHistoryItem(setHistory, { type: 'error', text: `[Error starting stream: ${error.message}]` }, getNextMessageId(userMessageTimestamp)); - } - } finally { - abortControllerRef.current = null; - setStreamingState(StreamingState.Idle); + // Call the renamed processor function + await processGeminiStream({ + stream, + signal, + setHistory, + submitQuery, + getNextMessageId: getStreamMessageId, + addHistoryItem: addHistoryItemFromStream, + currentToolGroupIdRef, + }); + } catch (error: any) { + // (Error handling for stream initiation remains the same) + console.error('Error initiating stream:', error); + if (error.name !== 'AbortError') { + // Use historyUpdater's function potentially? Or keep addHistoryItem here? + // Keeping addHistoryItem here for direct errors from this scope. + addHistoryItem( + setHistory, + { + type: 'error', + text: `[Error starting stream: ${error.message}]`, + }, + getNextMessageId(userMessageTimestamp), + ); } - }, [setStreamingState, setHistory, initError, getNextMessageId]); + } finally { + abortControllerRef.current = null; + setStreamingState(StreamingState.Idle); + } + }, + [setStreamingState, setHistory, initError, getNextMessageId], + ); - return { streamingState, submitQuery, initError }; + return { streamingState, submitQuery, initError }; }; diff --git a/packages/cli/src/ui/hooks/useLoadingIndicator.ts b/packages/cli/src/ui/hooks/useLoadingIndicator.ts index 8f440327..7641cd01 100644 --- a/packages/cli/src/ui/hooks/useLoadingIndicator.ts +++ b/packages/cli/src/ui/hooks/useLoadingIndicator.ts @@ -1,53 +1,61 @@ import { useState, useEffect, useRef } from 'react'; -import { WITTY_LOADING_PHRASES, PHRASE_CHANGE_INTERVAL_MS } from '../constants.js'; +import { + WITTY_LOADING_PHRASES, + PHRASE_CHANGE_INTERVAL_MS, +} from '../constants.js'; import { StreamingState } from '../../core/gemini-stream.js'; export const useLoadingIndicator = (streamingState: StreamingState) => { - const [elapsedTime, setElapsedTime] = useState(0); - const [currentLoadingPhrase, setCurrentLoadingPhrase] = useState(WITTY_LOADING_PHRASES[0]); - const timerRef = useRef<NodeJS.Timeout | null>(null); - const phraseIntervalRef = useRef<NodeJS.Timeout | null>(null); - const currentPhraseIndexRef = useRef<number>(0); + const [elapsedTime, setElapsedTime] = useState(0); + const [currentLoadingPhrase, setCurrentLoadingPhrase] = useState( + WITTY_LOADING_PHRASES[0], + ); + const timerRef = useRef<NodeJS.Timeout | null>(null); + const phraseIntervalRef = useRef<NodeJS.Timeout | null>(null); + const currentPhraseIndexRef = useRef<number>(0); - // Timer effect for elapsed time during loading - useEffect(() => { - if (streamingState === StreamingState.Responding) { - setElapsedTime(0); // Reset timer on new loading start - timerRef.current = setInterval(() => { - setElapsedTime((prevTime) => prevTime + 1); - }, 1000); - } else if (timerRef.current) { - clearInterval(timerRef.current); - timerRef.current = null; - } - // Cleanup on unmount or when isLoading changes - return () => { - if (timerRef.current) { - clearInterval(timerRef.current); - } - }; - }, [streamingState]); + // Timer effect for elapsed time during loading + useEffect(() => { + if (streamingState === StreamingState.Responding) { + setElapsedTime(0); // Reset timer on new loading start + timerRef.current = setInterval(() => { + setElapsedTime((prevTime) => prevTime + 1); + }, 1000); + } else if (timerRef.current) { + clearInterval(timerRef.current); + timerRef.current = null; + } + // Cleanup on unmount or when isLoading changes + return () => { + if (timerRef.current) { + clearInterval(timerRef.current); + } + }; + }, [streamingState]); - // Effect for cycling through witty loading phrases - useEffect(() => { - if (streamingState === StreamingState.Responding) { - currentPhraseIndexRef.current = 0; - setCurrentLoadingPhrase(WITTY_LOADING_PHRASES[0]); - phraseIntervalRef.current = setInterval(() => { - currentPhraseIndexRef.current = (currentPhraseIndexRef.current + 1) % WITTY_LOADING_PHRASES.length; - setCurrentLoadingPhrase(WITTY_LOADING_PHRASES[currentPhraseIndexRef.current]); - }, PHRASE_CHANGE_INTERVAL_MS); - } else if (phraseIntervalRef.current) { - clearInterval(phraseIntervalRef.current); - phraseIntervalRef.current = null; - } - // Cleanup on unmount or when isLoading changes - return () => { - if (phraseIntervalRef.current) { - clearInterval(phraseIntervalRef.current); - } - }; - }, [streamingState]); + // Effect for cycling through witty loading phrases + useEffect(() => { + if (streamingState === StreamingState.Responding) { + currentPhraseIndexRef.current = 0; + setCurrentLoadingPhrase(WITTY_LOADING_PHRASES[0]); + phraseIntervalRef.current = setInterval(() => { + currentPhraseIndexRef.current = + (currentPhraseIndexRef.current + 1) % WITTY_LOADING_PHRASES.length; + setCurrentLoadingPhrase( + WITTY_LOADING_PHRASES[currentPhraseIndexRef.current], + ); + }, PHRASE_CHANGE_INTERVAL_MS); + } else if (phraseIntervalRef.current) { + clearInterval(phraseIntervalRef.current); + phraseIntervalRef.current = null; + } + // Cleanup on unmount or when isLoading changes + return () => { + if (phraseIntervalRef.current) { + clearInterval(phraseIntervalRef.current); + } + }; + }, [streamingState]); - return { elapsedTime, currentLoadingPhrase }; -};
\ No newline at end of file + return { elapsedTime, currentLoadingPhrase }; +}; |
