update playground loading to support simulated stream

latitude-dev · Oct 18, 2024 · 2f9bc12 · 2f9bc12
1 parent 77a4878
commit 2f9bc12
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 10 deletions.
diff --git a/...on)/[evaluationUuid]/editor/_components/EvaluationEditor/Editor/Playground/Chat/index.tsx b/...on)/[evaluationUuid]/editor/_components/EvaluationEditor/Editor/Playground/Chat/index.tsx
@@ -50,6 +50,7 @@ export default function Chat({
   const [chainLength, setChainLength] = useState<number>(Infinity)
   const [conversation, setConversation] = useState<Conversation | undefined>()
   const [responseStream, setResponseStream] = useState<string | undefined>()
+  const [isStreaming, setIsStreaming] = useState(false)
 
   const addMessageToConversation = useCallback(
     (message: ConversationMessage) => {
@@ -69,7 +70,7 @@ export default function Chat({
   const runEvaluation = useCallback(async () => {
     setError(undefined)
     setResponseStream(undefined)
-
+    setIsStreaming(true)
     let response = ''
     let messagesCount = 0
 
@@ -101,9 +102,11 @@ export default function Chat({
           } else if (data.type === ChainEventTypes.Complete) {
             setResponseStream(undefined)
             setUsage(data.response.usage)
+            setIsStreaming(false)
             setEndTime(performance.now())
           } else if (data.type === ChainEventTypes.Error) {
             setError(new Error(data.error.message))
+            setIsStreaming(false)
           }
           break
         }
@@ -119,6 +122,7 @@ export default function Chat({
           break
       }
     }
+    setIsStreaming(false)
   }, [parameters, runPromptAction])
 
   useEffect(() => {
@@ -168,11 +172,16 @@ export default function Chat({
         <TokenUsage
           isScrolledToBottom={isScrolledToBottom}
           usage={usage}
-          responseStream={responseStream}
+          isStreaming={isStreaming}
         />
       </div>
       <div className='flex items-center justify-center'>
-        <Button fancy variant='outline' onClick={clearChat}>
+        <Button
+          disabled={isStreaming}
+          fancy
+          variant='outline'
+          onClick={clearChat}
+        >
           Clear chat
         </Button>
       </div>

diff --git a/...ommitUuid]/documents/[documentUuid]/_components/DocumentEditor/Editor/Playground/Chat.tsx b/...ommitUuid]/documents/[documentUuid]/_components/DocumentEditor/Editor/Playground/Chat.tsx
@@ -60,6 +60,7 @@ export default function Chat({
   const [chainLength, setChainLength] = useState<number>(Infinity)
   const [conversation, setConversation] = useState<Conversation | undefined>()
   const [responseStream, setResponseStream] = useState<string | undefined>()
+  const [isStreaming, setIsStreaming] = useState(false)
 
   const addMessageToConversation = useCallback(
     (message: ConversationMessage) => {
@@ -80,7 +81,12 @@ export default function Chat({
     const start = performance.now()
     setError(undefined)
     setResponseStream('')
-
+    setIsStreaming(true)
+    setUsage({
+      promptTokens: 0,
+      completionTokens: 0,
+      totalTokens: 0,
+    })
     let response = ''
     let messagesCount = 0
 
@@ -136,6 +142,7 @@ export default function Chat({
     } catch (error) {
       setError(error as Error)
     } finally {
+      setIsStreaming(false)
       setResponseStream(undefined)
     }
   }, [
@@ -258,12 +265,12 @@ export default function Chat({
         <TokenUsage
           isScrolledToBottom={isScrolledToBottom}
           usage={usage}
-          responseStream={responseStream}
+          isStreaming={isStreaming}
         />
         <ChatTextArea
           clearChat={clearChat}
           placeholder='Enter followup message...'
-          disabled={responseStream !== undefined}
+          disabled={isStreaming}
           onSubmit={submitUserMessage}
         />
       </div>
@@ -290,13 +297,13 @@ export function AnimatedDots() {
 export function TokenUsage({
   isScrolledToBottom,
   usage,
-  responseStream,
+  isStreaming,
 }: {
   isScrolledToBottom: boolean
   usage: LanguageModelUsage | undefined
-  responseStream: string | undefined
+  isStreaming: boolean
 }) {
-  if (!usage && responseStream === undefined) return null
+  if (!usage && isStreaming) return null
 
   return (
     <div
@@ -307,7 +314,7 @@ export function TokenUsage({
         },
       )}
     >
-      {responseStream === undefined ? (
+      {!isStreaming && usage ? (
         <Tooltip
           side='top'
           align='center'