Refactor UI to use local state for pipeline / status change tracking

Refactor LitLytics class to always return new values on actions Move step addition logic to litlytics class
yamalight · Oct 16, 2024 · 4b24453 · 4b24453
1 parent f2eafba
commit 4b24453
Show file tree

Hide file tree

Showing 16 changed files with 413 additions and 285 deletions.
diff --git a/app/components/pipeline/GeneratePipeline.tsx b/app/components/pipeline/GeneratePipeline.tsx
@@ -24,7 +24,7 @@ const tabClass = clsx(
 );
 
 export default function GeneratePipeline() {
-  const litlytics = useLitlytics();
+  const { litlytics, pipeline, setPipeline, pipelineStatus } = useLitlytics();
   const [selectedTab, setSelectedTab] = useState(0);
   const [isOpen, setIsOpen] = useState(false);
   const [loading, setLoading] = useState(false);
@@ -36,8 +36,8 @@ export default function GeneratePipeline() {
       setError(undefined);
 
       // generate plan from LLM
-      await litlytics.generatePipeline();
-
+      const newPipeline = await litlytics.generatePipeline();
+      setPipeline(newPipeline);
       setLoading(false);
       setSelectedTab(1);
     } catch (err) {
@@ -49,9 +49,9 @@ export default function GeneratePipeline() {
   const closeDialog = () => {
     if (
       loading ||
-      litlytics.pipelineStatus.status === 'refine' ||
-      litlytics.pipelineStatus.status === 'step' ||
-      litlytics.pipelineStatus.status === 'sourcing'
+      pipelineStatus.status === 'refine' ||
+      pipelineStatus.status === 'step' ||
+      pipelineStatus.status === 'sourcing'
     ) {
       return;
     }
@@ -65,8 +65,8 @@ export default function GeneratePipeline() {
         outline
         onClick={() => setIsOpen(true)}
         disabled={
-          litlytics.pipelineStatus.status === 'sourcing' ||
-          litlytics.pipelineStatus.status === 'step'
+          pipelineStatus.status === 'sourcing' ||
+          pipelineStatus.status === 'step'
         }
         className="border-sky-500 dark:border-sky-500"
       >
@@ -83,7 +83,7 @@ export default function GeneratePipeline() {
           <TabGroup selectedIndex={selectedTab} onChange={setSelectedTab}>
             <TabList className="flex gap-4">
               <Tab className={tabClass}>Plan</Tab>
-              {Boolean(litlytics.pipeline.pipelinePlan?.length) && (
+              {Boolean(pipeline.pipelinePlan?.length) && (
                 <Tab className={tabClass}>Refine plan</Tab>
               )}
             </TabList>
@@ -96,9 +96,10 @@ export default function GeneratePipeline() {
                       rows={5}
                       name="task"
                       placeholder="Task description"
-                      value={litlytics.pipeline.pipelineDescription}
+                      value={pipeline.pipelineDescription ?? ''}
                       onChange={(e) => {
-                        litlytics.setPipeline({
+                        setPipeline({
+                          ...pipeline,
                           pipelineDescription: e.target.value,
                         });
                       }}

diff --git a/app/components/pipeline/PipelineBuilder.tsx b/app/components/pipeline/PipelineBuilder.tsx
@@ -5,12 +5,12 @@ import { SourceNode } from './nodes/source/SourceNode';
 import { StepNode } from './nodes/StepNode';
 
 export function PipelineBuilder() {
-  const litlytics = useLitlytics();
+  const { pipeline } = useLitlytics();
 
   return (
     <Background>
       <SourceNode />
-      {litlytics.pipeline.steps
+      {pipeline.steps
         .sort((a, b) => (a.connectsTo.includes(b.id) ? -1 : 1))
         .map((step) => (
           <StepNode data={step} key={step.id} />

diff --git a/app/components/pipeline/RefinePipeline.tsx b/app/components/pipeline/RefinePipeline.tsx
@@ -6,7 +6,13 @@ import { Spinner } from '~/components/Spinner';
 import { useLitlytics } from '~/store/WithLitLytics';
 
 export function RefinePipeline({ hide }: { hide: () => void }) {
-  const litlytics = useLitlytics();
+  const {
+    litlytics,
+    pipeline,
+    setPipeline,
+    pipelineStatus,
+    setPipelineStatus,
+  } = useLitlytics();
   const [error, setError] = useState<Error>();
   const [refine, setRefine] = useState(``);
   const [progress, setProgress] = useState('');
@@ -18,48 +24,54 @@ export function RefinePipeline({ hide }: { hide: () => void }) {
 
     try {
       // generate plan from LLM
-      litlytics.setPipelineStatus({ status: 'refine' });
-      await litlytics.refinePipeline({
+      setPipelineStatus({ status: 'refine' });
+      const newPipeline = await litlytics.refinePipeline({
         refineRequest: refine,
       });
-      litlytics.setPipelineStatus({ status: 'init' });
+      setPipeline(newPipeline);
+      setPipelineStatus({ status: 'init' });
       setRefine('');
     } catch (err) {
       setError(err as Error);
-      litlytics.setPipelineStatus({ status: 'error' });
+      setPipelineStatus({ status: 'error' });
     }
   };
 
   const doCreate = async () => {
     try {
+      setPipelineStatus({ status: 'sourcing' });
       // generate plan from LLM
-      await litlytics.pipelineFromText(({ step, totalSteps }) => {
-        if (step > totalSteps) {
-          setProgress('');
-          return;
-        }
+      const newPipeline = await litlytics.pipelineFromText(
+        ({ step, totalSteps }) => {
+          if (step > totalSteps) {
+            setProgress('');
+            return;
+          }
 
-        setProgress(`Generating steps: ${step} / ${totalSteps}`);
-      });
+          setProgress(`Generating steps: ${step} / ${totalSteps}`);
+        }
+      );
+      setPipeline(newPipeline);
+      setPipelineStatus({ status: 'done' });
       hide();
     } catch (err) {
       setError(err as Error);
-      litlytics.setPipelineStatus({ status: 'error' });
+      setPipelineStatus({ status: 'error' });
     }
   };
 
   const inProgress =
-    litlytics.pipelineStatus.status === 'refine' ||
-    litlytics.pipelineStatus.status === 'sourcing' ||
-    litlytics.pipelineStatus.status === 'step';
+    pipelineStatus.status === 'refine' ||
+    pipelineStatus.status === 'sourcing' ||
+    pipelineStatus.status === 'step';
 
   return (
     <div className="w-full h-full p-4 prose prose-sm prose-no-nr dark:prose-invert ">
       <div className="flex w-full items-center justify-between">
         <h1 className="m-0">Suggested pipeline:</h1>
       </div>
 
-      <CustomMarkdown>{`${litlytics.pipeline.pipelinePlan}`}</CustomMarkdown>
+      <CustomMarkdown>{`${pipeline.pipelinePlan}`}</CustomMarkdown>
       <div className="flex gap-1">
         <Textarea
           rows={2}
@@ -69,7 +81,7 @@ export function RefinePipeline({ hide }: { hide: () => void }) {
           onChange={(e) => setRefine(e.target.value)}
         />
         <Button onClick={doRefine} disabled={inProgress}>
-          {litlytics.pipelineStatus.status === 'refine' && (
+          {pipelineStatus.status === 'refine' && (
             <Spinner className="h-5 w-5" />
           )}
           Refine
@@ -84,14 +96,14 @@ export function RefinePipeline({ hide }: { hide: () => void }) {
 
       <div className="flex mt-8">
         <Button onClick={doCreate} disabled={inProgress}>
-          {(litlytics.pipelineStatus.status === 'sourcing' ||
-            litlytics.pipelineStatus.status === 'step') && (
+          {(pipelineStatus.status === 'sourcing' ||
+            pipelineStatus.status === 'step') && (
             <div className="flex items-center">
               <Spinner className="h-5 w-5" />
             </div>
           )}
-          {litlytics.pipelineStatus.status === 'sourcing' ||
-          litlytics.pipelineStatus.status === 'step'
+          {pipelineStatus.status === 'sourcing' ||
+          pipelineStatus.status === 'step'
             ? progress.length > 0
               ? progress
               : `Generating pipeline...`

diff --git a/app/components/pipeline/nodes/NodeConnector.tsx b/app/components/pipeline/nodes/NodeConnector.tsx
@@ -1,7 +1,6 @@
 import { PlusIcon } from '@heroicons/react/24/solid';
 import clsx from 'clsx';
 import {
-  OUTPUT_ID,
   ProcessingStep,
   ProcessingStepTypes,
   SourceStep,
@@ -47,7 +46,7 @@ export function NodeConnector({
   currentStep?: SourceStep | ProcessingStep;
   showAuto?: boolean;
 }) {
-  const litlytics = useLitlytics();
+  const { litlytics, setPipeline } = useLitlytics();
   const [step, setStep] = useState<ProcessingStep>(
     structuredClone(defaultStep)
   );
@@ -66,68 +65,13 @@ export function NodeConnector({
     }
 
     setLoading(true);
-    // generate new ID and double-check that it doesn't overlap with other steps
-    let id = litlytics.pipeline.steps.length;
-    let existingStep = litlytics.pipeline.steps.find(
-      (s) => s.id === `step_${id}`
-    );
-    while (existingStep) {
-      id += 1;
-      existingStep = litlytics.pipeline.steps.find(
-        (s) => s.id === `step_${id}`
-      );
-    }
-    // generate final id
-    const idStr = `step_${id}`;
-
-    let newStep: ProcessingStep | undefined = undefined;
-    if (manual) {
-      newStep = structuredClone(step);
-      newStep.id = idStr;
-    } else {
-      // generate new step
-      newStep = await litlytics.generateStep({
-        id: idStr,
-        name: step.name,
-        description: step.description,
-        input: step.input as StepInput,
-        type: step.type,
-      });
-    }
 
-    if (currentStep?.type === 'source') {
-      // connect new step to next node
-      const nextNodeId =
-        litlytics.pipeline.source.connectsTo.at(0) ?? OUTPUT_ID;
-      newStep.connectsTo = [nextNodeId];
-      // add
-      litlytics.setPipeline({
-        source: {
-          ...litlytics.pipeline.source,
-          connectsTo: [newStep.id],
-        },
-        steps: litlytics.pipeline.steps.concat(newStep),
-      });
-    } else {
-      // connect new step to next node
-      const nextNodeId =
-        litlytics.pipeline.steps
-          .find((s) => s.id === currentStep?.id)
-          ?.connectsTo.at(0) ?? OUTPUT_ID;
-      newStep.connectsTo = [nextNodeId];
-      // add
-      litlytics.setPipeline({
-        steps: litlytics.pipeline.steps
-          .map((s) => {
-            if (s.id === currentStep?.id) {
-              s.connectsTo = [newStep.id];
-              return s;
-            }
-            return s;
-          })
-          .concat(newStep),
-      });
-    }
+    const newPipeline = await litlytics.addStep({
+      step,
+      sourceStep: currentStep,
+      manual,
+    });
+    setPipeline(newPipeline);
 
     setIsOpen(false);
     setLoading(false);