feat: add trace stream toggle into preferences context (#4035)

Arize-ai · Jul 27, 2024 · bc3be3e · bc3be3e
1 parent 09a5765
commit bc3be3e
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 3 deletions.
diff --git a/app/src/components/markdown/MarkdownDisplayContext.tsx b/app/src/components/markdown/MarkdownDisplayContext.tsx
@@ -31,7 +31,9 @@ export function useMarkdownMode(): MarkdownDisplayContextType {
 }
 
 export function MarkdownDisplayProvider(props: PropsWithChildren) {
-  const mode = usePreferencesContext((state) => state.markdownDisplayMode);
+  const mode = usePreferencesContext((state) => {
+    return state.markdownDisplayMode;
+  });
   const setMarkdownDisplayMode = usePreferencesContext(
     (state) => state.setMarkdownDisplayMode
   );

diff --git a/app/src/contexts/StreamStateContext.tsx b/app/src/contexts/StreamStateContext.tsx
@@ -6,6 +6,8 @@ import React, {
   useState,
 } from "react";
 
+import { usePreferencesContext } from "./PreferencesContext";
+
 export type StreamStateContextType = {
   /**
    * Whether or not streaming is enabled.
@@ -44,7 +46,12 @@ export function StreamStateProvider({
   initialFetchKey = "initial",
   children,
 }: PropsWithChildren<{ initialFetchKey?: string }>) {
-  const [isStreaming, setIsStreaming] = useState<boolean>(true);
+  const isStreaming = usePreferencesContext(
+    (state) => state.traceStreamingEnabled
+  );
+  const setIsStreaming = usePreferencesContext(
+    (state) => state.setTraceStreamingEnabled
+  );
   const [fetchKey, _setFetchKey] = useState<string>(initialFetchKey);
 
   const setFetchKey = useCallback(

diff --git a/app/src/pages/project/StreamToggle.tsx b/app/src/pages/project/StreamToggle.tsx
@@ -58,7 +58,7 @@ export function StreamToggle(props: { project: StreamToggle_data$key }) {
   return (
     <Switch
       labelPlacement="start"
-      defaultSelected
+      isSelected={isStreaming}
       onChange={() => {
         setIsStreaming(!isStreaming);
       }}

diff --git a/app/src/store/preferencesStore.tsx b/app/src/store/preferencesStore.tsx
@@ -17,6 +17,17 @@ export interface PreferencesState extends PreferencesProps {
    * @param markdownDisplayMode
    */
   setMarkdownDisplayMode: (markdownDisplayMode: MarkdownDisplayMode) => void;
+  /**
+   * Whether or not streaming is enabled for a projects traces
+   * @default true
+   */
+  traceStreamingEnabled: boolean;
+  /**
+   * Setter for enabling/disabling trace streaming
+   * @param traceStreamingEnabled
+   * @returns
+   */
+  setTraceStreamingEnabled: (traceStreamingEnabled: boolean) => void;
 }
 
 export const createPreferencesStore = (
@@ -28,6 +39,10 @@ export const createPreferencesStore = (
     setMarkdownDisplayMode: (markdownDisplayMode) => {
       set({ markdownDisplayMode });
     },
+    traceStreamingEnabled: true,
+    setTraceStreamingEnabled: (traceStreamingEnabled) => {
+      set({ traceStreamingEnabled });
+    },
   });
   return create<PreferencesState>()(
     persist(devtools(preferencesStore), {