Beam: Persist (and get off the way) more state

This commit is contained in:
Enrico Ros
2024-03-23 14:46:15 -07:00
parent 79a8fbd881
commit 5bfcef92ee
7 changed files with 61 additions and 71 deletions
+2 -3
View File
@@ -34,11 +34,11 @@ export function BeamView(props: {
const {
/* root */ inputHistory, inputIssues, inputReady,
/* scatter */ isScattering, raysReady,
/* gather */ gatherLlmId, gatherShowPrompts,
/* gather */ gatherLlmId,
} = useBeamStore(props.beamStore, useShallow(state => ({
inputHistory: state.inputHistory, inputIssues: state.inputIssues, inputReady: state.inputReady,
isScattering: state.isScattering, raysReady: state.raysReady,
gatherLlmId: state.gatherLlmId, gatherShowPrompts: state.gatherShowPrompts,
gatherLlmId: state.gatherLlmId,
})));
const rayIds = useBeamStore(props.beamStore, useShallow(state => state.rays.map(ray => ray.rayId)));
const [_, gatherLlmComponent, gatherLlmIcon] = useLLMSelect(gatherLlmId, setGatherLlmId, props.isMobile ? '' : 'Merge Model', true);
@@ -129,7 +129,6 @@ export function BeamView(props: {
{/* Fusion Config */}
<BeamGatherInput
beamStore={props.beamStore}
gatherShowPrompts={gatherShowPrompts}
/>
{/* Gather Controls */}
+7 -4
View File
@@ -11,6 +11,7 @@ import type { ChatGenerateInstruction, Instruction } from './beam.gather.instruc
import { BeamStoreApi, useBeamStore } from '../store-beam.hooks';
import { GATHER_SHOW_SYSTEM_PROMPT } from '../beam.config';
import { fusionIsEditable } from './beam.gather';
import { useModuleBeamStore } from '../store-module-beam';
const gatherInputWrapperSx: SxProps = {
@@ -175,11 +176,13 @@ function EditableInstruction(props: {
export function BeamGatherInput(props: {
beamStore: BeamStoreApi,
gatherShowPrompts: boolean
}) {
// external state (all null if we don't have an index)
const { currentFusionId, currentIsEditable, currentInstructions } = useBeamStore(props.beamStore, useShallow(store => {
// external state
const gatherShowPrompts = useModuleBeamStore(state => state.gatherShowPrompts);
const {
currentFusionId, currentIsEditable, currentInstructions,
} = useBeamStore(props.beamStore, useShallow(store => {
const fusion = store.currentFusionId !== null ? store.fusions.find(fusion => fusion.fusionId === store.currentFusionId) ?? null : null;
return {
currentFusionId: fusion?.fusionId ?? null,
@@ -230,7 +233,7 @@ export function BeamGatherInput(props: {
// render if existing and editable
if (currentFusionId === null || (!currentIsEditable && !props.gatherShowPrompts))
if (currentFusionId === null || (!currentIsEditable && !gatherShowPrompts))
return null;
return (
+17 -33
View File
@@ -8,8 +8,6 @@ import AutoAwesomeOutlinedIcon from '@mui/icons-material/AutoAwesomeOutlined';
import MergeRoundedIcon from '@mui/icons-material/MergeRounded';
import StopRoundedIcon from '@mui/icons-material/StopRounded';
import { BeamGatherDropdown } from '~/common/beam/gather/BeamGatherPaneDropdown';
import { BeamStoreApi, useBeamStore } from '~/common/beam/store-beam.hooks';
import { ConfirmationModal } from '~/common/components/ConfirmationModal';
import { FormLabelStart } from '~/common/components/forms/FormLabelStart';
import { GoodTooltip } from '~/common/components/GoodTooltip';
@@ -17,9 +15,12 @@ import { ScrollToBottomButton } from '~/common/scroll-to-bottom/ScrollToBottomBu
import { animationColorBeamGather, animationShadowLimey } from '~/common/util/animUtils';
import { useScrollToBottom } from '~/common/scroll-to-bottom/useScrollToBottom';
import { BeamGatherDropdown } from './BeamGatherPaneDropdown';
import { BeamStoreApi, useBeamStore } from '../store-beam.hooks';
import { FUSION_FACTORIES } from './beam.gather.factories';
import { GATHER_COLOR } from '../beam.config';
import { beamPaneSx } from '../BeamCard';
import { useModuleBeamStore } from '../store-module-beam';
const gatherPaneClasses = {
@@ -73,31 +74,24 @@ export function BeamGatherPane(props: {
// external state
const {
gatherShowDevMethods, gatherShowPrompts,
fusions, currentFusionId, isGatheringAny,
toggleGatherShowDevMethods, toggleGatherShowPrompts,
setCurrentFusionId, currentFusionStart, currentFusionStop,
stopScatteringAll,
} = useBeamStore(props.beamStore, useShallow(state => {
return {
// state (gatherLlmId is lifted to the parent)
gatherShowDevMethods: state.gatherShowDevMethods,
gatherShowPrompts: state.gatherShowPrompts,
currentFusionId: state.currentFusionId,
fusions: state.fusions,
isGatheringAny: state.isGatheringAny,
} = useBeamStore(props.beamStore, useShallow(state => ({
// state (gatherLlmId is lifted to the parent)
currentFusionId: state.currentFusionId,
fusions: state.fusions,
isGatheringAny: state.isGatheringAny,
// actions
toggleGatherShowDevMethods: state.toggleGatherShowDevMethods,
toggleGatherShowPrompts: state.toggleGatherShowPrompts,
setCurrentFusionId: state.setCurrentFusionId,
currentFusionStart: state.currentFusionStart,
currentFusionStop: state.currentFusionStop,
// actions
setCurrentFusionId: state.setCurrentFusionId,
currentFusionStart: state.currentFusionStart,
currentFusionStop: state.currentFusionStop,
// (external slice) scatter actions
stopScatteringAll: state.stopScatteringAll,
};
}));
// (external slice) scatter actions
stopScatteringAll: state.stopScatteringAll,
})));
const gatherShowDevMethods = useModuleBeamStore(state => state.gatherShowDevMethods);
const { setStickToBottom } = useScrollToBottom();
@@ -145,16 +139,6 @@ export function BeamGatherPane(props: {
}, [handleStopScatterConfirmation, props.scatterBusy, warnScatterBusy]);
const dropdownMemo = React.useMemo(() => (
<BeamGatherDropdown
gatherShowDevMethods={gatherShowDevMethods}
gatherShowPrompts={gatherShowPrompts}
toggleGatherShowDevMethods={toggleGatherShowDevMethods}
toggleGatherShowPrompts={toggleGatherShowPrompts}
/>
), [gatherShowDevMethods, gatherShowPrompts, toggleGatherShowDevMethods, toggleGatherShowPrompts]);
const MainLlmIcon = props.gatherLlmIcon || (isGatheringAny ? AutoAwesomeIcon : AutoAwesomeOutlinedIcon);
return <>
@@ -168,7 +152,7 @@ export function BeamGatherPane(props: {
<div>
<Typography
level='h4' component='h2'
endDecorator={dropdownMemo}
endDecorator={<BeamGatherDropdown />}
// sx={{ my: 0.25 }}
>
<MainLlmIcon sx={{ fontSize: '1rem', animation: isGatheringAny ? `${animationColorBeamGather} 2s linear infinite` : undefined }} />&nbsp;Merge
@@ -1,16 +1,28 @@
import * as React from 'react';
import { useShallow } from 'zustand/react/shallow';
import { Dropdown, IconButton, ListItem, ListItemDecorator, Menu, MenuButton, MenuItem, Typography } from '@mui/joy';
import CheckRoundedIcon from '@mui/icons-material/CheckRounded';
import MoreHorizRoundedIcon from '@mui/icons-material/MoreHorizRounded';
import { useModuleBeamStore } from '../store-module-beam';
export function BeamGatherDropdown() {
// external (persisted) state
const {
gatherShowDevMethods,
gatherShowPrompts,
toggleGatherShowDevMethods,
toggleGatherShowPrompts,
} = useModuleBeamStore(useShallow(state => ({
gatherShowDevMethods: state.gatherShowDevMethods,
gatherShowPrompts: state.gatherShowPrompts,
toggleGatherShowDevMethods: state.toggleGatherShowDevMethods,
toggleGatherShowPrompts: state.toggleGatherShowPrompts,
})));
export function BeamGatherDropdown(props: {
gatherShowDevMethods: boolean,
gatherShowPrompts: boolean,
toggleGatherShowDevMethods: () => void,
toggleGatherShowPrompts: () => void,
}) {
return (
<Dropdown>
<MenuButton
@@ -24,12 +36,12 @@ export function BeamGatherDropdown(props: {
<ListItem>
<Typography level='body-sm'>Advanced</Typography>
</ListItem>
<MenuItem onClick={props.toggleGatherShowPrompts}>
<ListItemDecorator>{props.gatherShowPrompts && <CheckRoundedIcon />}</ListItemDecorator>
<MenuItem onClick={toggleGatherShowPrompts}>
<ListItemDecorator>{gatherShowPrompts && <CheckRoundedIcon />}</ListItemDecorator>
Show All Prompts
</MenuItem>
<MenuItem onClick={props.toggleGatherShowDevMethods}>
<ListItemDecorator>{props.gatherShowDevMethods && <CheckRoundedIcon />}</ListItemDecorator>
<MenuItem onClick={toggleGatherShowDevMethods}>
<ListItemDecorator>{gatherShowDevMethods && <CheckRoundedIcon />}</ListItemDecorator>
Experimental Methods
</MenuItem>
</Menu>
-18
View File
@@ -65,9 +65,6 @@ interface GatherStateSlice {
gatherLlmId: DLLMId | null;
gatherShowDevMethods: boolean;
gatherShowPrompts: boolean;
currentFusionId: BFusionId | null;
fusions: BFusion[];
@@ -86,9 +83,6 @@ export const reInitGatherStateSlice = (prevFusions: BFusion[]): GatherStateSlice
return {
gatherLlmId: null, // will be re-set during open() of the Beam Store
gatherShowDevMethods: false,
gatherShowPrompts: false,
currentFusionId: (GATHER_DEFAULT_TO_FIRST_FUSION && newFusions.length) ? newFusions[0].fusionId : null,
fusions: newFusions,
@@ -101,8 +95,6 @@ export type FusionUpdateOrFn = Partial<BFusion> | ((fusion: BFusion) => (Partial
export interface GatherStoreSlice extends GatherStateSlice {
setGatherLlmId: (llmId: DLLMId | null) => void;
toggleGatherShowDevMethods: () => void;
toggleGatherShowPrompts: () => void;
setCurrentFusionId: (fusionId: BFusionId | null) => void;
_currentFusion: () => BFusion | null;
@@ -127,16 +119,6 @@ export const createGatherSlice: StateCreator<GatherStoreSlice, [], [], GatherSto
gatherLlmId: llmId,
}),
toggleGatherShowDevMethods: () =>
_set(state => ({
gatherShowDevMethods: !state.gatherShowDevMethods,
})),
toggleGatherShowPrompts: () =>
_set(state => ({
gatherShowPrompts: !state.gatherShowPrompts,
})),
setCurrentFusionId: (fusionId: BFusionId | null) =>
_set({
+1 -3
View File
@@ -94,7 +94,7 @@ const createRootSlice: StateCreator<BeamStore, [], [], RootStoreSlice> = (_set,
},
terminate: () => { /*_get().isOpen &&*/
const { rays, fusions, gatherLlmId, gatherShowDevMethods, gatherShowPrompts } = _get();
const { rays, fusions, gatherLlmId } = _get();
_set({
...initRootStateSlice(),
@@ -103,8 +103,6 @@ const createRootSlice: StateCreator<BeamStore, [], [], RootStoreSlice> = (_set,
// remember after termination
gatherLlmId,
gatherShowDevMethods,
gatherShowPrompts,
});
},
+12
View File
@@ -18,6 +18,8 @@ interface ModuleBeamStore {
// state
scatterPresets: BeamScatterPreset[];
rayScrolling: boolean;
gatherShowDevMethods: boolean;
gatherShowPrompts: boolean;
// actions
addScatterPreset: (name: string, scatterLlmIds: DLLMId[]) => void;
@@ -25,6 +27,9 @@ interface ModuleBeamStore {
renameScatterPreset: (id: string, name: string) => void;
toggleRayScrolling: () => void;
toggleGatherShowDevMethods: () => void;
toggleGatherShowPrompts: () => void;
}
@@ -33,6 +38,8 @@ export const useModuleBeamStore = create<ModuleBeamStore>()(persist(
scatterPresets: [],
rayScrolling: false,
gatherShowDevMethods: false,
gatherShowPrompts: false,
addScatterPreset: (name, scatterLlmIds) => _set(state => ({
@@ -50,6 +57,11 @@ export const useModuleBeamStore = create<ModuleBeamStore>()(persist(
toggleRayScrolling: () => _set(state => ({ rayScrolling: !state.rayScrolling })),
toggleGatherShowDevMethods: () => _set(state => ({ gatherShowDevMethods: !state.gatherShowDevMethods })),
toggleGatherShowPrompts: () => _set(state => ({ gatherShowPrompts: !state.gatherShowPrompts })),
}), {
name: 'app-module-beam',
},