Add new tasks suggestions flow to lib and UI

UI can use some extra work, but it's mostly OK for now

Author: yamalight

app/components/catalyst/listbox.tsx

@@ -0,0 +1,218 @@
+import * as Headless from '@headlessui/react';
+import clsx from 'clsx';
+import { Fragment } from 'react';
+
+export function Listbox<T>({
+  className,
+  optionsClassName,
+  placeholder,
+  autoFocus,
+  'aria-label': ariaLabel,
+  children: options,
+  ...props
+}: {
+  className?: string;
+  optionsClassName?: string;
+  placeholder?: React.ReactNode;
+  autoFocus?: boolean;
+  'aria-label'?: string;
+  children?: React.ReactNode;
+} & Omit<Headless.ListboxProps<typeof Fragment, T>, 'as' | 'multiple'>) {
+  return (
+    <Headless.Listbox {...props} multiple={false}>
+      <Headless.ListboxButton
+        autoFocus={autoFocus}
+        data-slot="control"
+        aria-label={ariaLabel}
+        className={clsx([
+          className,
+          // Basic layout
+          'group relative block w-full',
+          // Background color + shadow applied to inset pseudo element, so shadow blends with border in light mode
+          'before:absolute before:inset-px before:rounded-[calc(theme(borderRadius.lg)-1px)] before:bg-white before:shadow',
+          // Background color is moved to control and shadow is removed in dark mode so hide `before` pseudo
+          'dark:before:hidden',
+          // Hide default focus styles
+          'focus:outline-none',
+          // Focus ring
+          'after:pointer-events-none after:absolute after:inset-0 after:rounded-lg after:ring-inset after:ring-transparent after:data-[focus]:ring-2 after:data-[focus]:ring-blue-500',
+          // Disabled state
+          'data-[disabled]:opacity-50 before:data-[disabled]:bg-zinc-950/5 before:data-[disabled]:shadow-none',
+        ])}
+      >
+        <Headless.ListboxSelectedOption
+          as="span"
+          options={options}
+          placeholder={
+            placeholder && (
+              <span className="block truncate text-zinc-500">
+                {placeholder}
+              </span>
+            )
+          }
+          className={clsx([
+            // Basic layout
+            'relative block w-full appearance-none rounded-lg py-[calc(theme(spacing[2.5])-1px)] sm:py-[calc(theme(spacing[1.5])-1px)]',
+            // Set minimum height for when no value is selected
+            'min-h-11 sm:min-h-9',
+            // Horizontal padding
+            'pl-[calc(theme(spacing[3.5])-1px)] pr-[calc(theme(spacing.7)-1px)] sm:pl-[calc(theme(spacing.3)-1px)]',
+            // Typography
+            'text-left text-base/6 text-zinc-950 placeholder:text-zinc-500 sm:text-sm/6 dark:text-white forced-colors:text-[CanvasText]',
+            // Border
+            'border border-zinc-950/10 group-data-[active]:border-zinc-950/20 group-data-[hover]:border-zinc-950/20 dark:border-white/10 dark:group-data-[active]:border-white/20 dark:group-data-[hover]:border-white/20',
+            // Background color
+            'bg-transparent dark:bg-white/5',
+            // Invalid state
+            'group-data-[invalid]:border-red-500 group-data-[invalid]:group-data-[hover]:border-red-500 group-data-[invalid]:dark:border-red-600 group-data-[invalid]:data-[hover]:dark:border-red-600',
+            // Disabled state
+            'group-data-[disabled]:border-zinc-950/20 group-data-[disabled]:opacity-100 group-data-[disabled]:dark:border-white/15 group-data-[disabled]:dark:bg-white/[2.5%] dark:data-[hover]:group-data-[disabled]:border-white/15',
+          ])}
+        />
+        <span className="pointer-events-none absolute inset-y-0 right-0 flex items-center pr-2">
+          <svg
+            className="size-5 stroke-zinc-500 group-data-[disabled]:stroke-zinc-600 sm:size-4 dark:stroke-zinc-400 forced-colors:stroke-[CanvasText]"
+            viewBox="0 0 16 16"
+            aria-hidden="true"
+            fill="none"
+          >
+            <path
+              d="M5.75 10.75L8 13L10.25 10.75"
+              strokeWidth={1.5}
+              strokeLinecap="round"
+              strokeLinejoin="round"
+            />
+            <path
+              d="M10.25 5.25L8 3L5.75 5.25"
+              strokeWidth={1.5}
+              strokeLinecap="round"
+              strokeLinejoin="round"
+            />
+          </svg>
+        </span>
+      </Headless.ListboxButton>
+      <Headless.ListboxOptions
+        transition
+        anchor="selection start"
+        className={clsx(
+          optionsClassName,
+          // Anchor positioning
+          '[--anchor-offset:-1.625rem] [--anchor-padding:theme(spacing.4)] sm:[--anchor-offset:-1.375rem]',
+          // Base styles
+          'isolate w-max min-w-[calc(var(--button-width)+1.75rem)] select-none scroll-py-1 rounded-xl p-1',
+          // Invisible border that is only visible in `forced-colors` mode for accessibility purposes
+          'outline outline-1 outline-transparent focus:outline-none',
+          // Handle scrolling when menu won't fit in viewport
+          'overflow-y-scroll overscroll-contain',
+          // Popover background
+          'bg-white/75 backdrop-blur-xl dark:bg-zinc-800/75',
+          // Shadows
+          'shadow-lg ring-1 ring-zinc-950/10 dark:ring-inset dark:ring-white/10',
+          // Transitions
+          'transition-opacity duration-100 ease-in data-[transition]:pointer-events-none data-[closed]:data-[leave]:opacity-0'
+        )}
+      >
+        {options}
+      </Headless.ListboxOptions>
+    </Headless.Listbox>
+  );
+}
+
+export function ListboxOption<T>({
+  children,
+  className,
+  ...props
+}: { className?: string; children?: React.ReactNode } & Omit<
+  Headless.ListboxOptionProps<'div', T>,
+  'as' | 'className'
+>) {
+  const sharedClasses = clsx(
+    // Base
+    'flex min-w-0 items-center',
+    // Icons
+    '[&>[data-slot=icon]]:size-5 [&>[data-slot=icon]]:shrink-0 sm:[&>[data-slot=icon]]:size-4',
+    '[&>[data-slot=icon]]:text-zinc-500 [&>[data-slot=icon]]:group-data-[focus]/option:text-white [&>[data-slot=icon]]:dark:text-zinc-400',
+    'forced-colors:[&>[data-slot=icon]]:text-[CanvasText] forced-colors:[&>[data-slot=icon]]:group-data-[focus]/option:text-[Canvas]',
+    // Avatars
+    '[&>[data-slot=avatar]]:-mx-0.5 [&>[data-slot=avatar]]:size-6 sm:[&>[data-slot=avatar]]:size-5'
+  );
+
+  return (
+    <Headless.ListboxOption as={Fragment} {...props}>
+      {({ selectedOption }) => {
+        if (selectedOption) {
+          return (
+            <div className={clsx(className, sharedClasses)}>{children}</div>
+          );
+        }
+
+        return (
+          <div
+            className={clsx(
+              // Basic layout
+              'group/option grid cursor-default grid-cols-[theme(spacing.5),1fr] items-baseline gap-x-2 rounded-lg py-2.5 pl-2 pr-3.5 sm:grid-cols-[theme(spacing.4),1fr] sm:py-1.5 sm:pl-1.5 sm:pr-3',
+              // Typography
+              'text-base/6 text-zinc-950 sm:text-sm/6 dark:text-white forced-colors:text-[CanvasText]',
+              // Focus
+              'outline-none data-[focus]:bg-blue-500 data-[focus]:text-white',
+              // Forced colors mode
+              'forced-color-adjust-none forced-colors:data-[focus]:bg-[Highlight] forced-colors:data-[focus]:text-[HighlightText]',
+              // Disabled
+              'data-[disabled]:opacity-50'
+            )}
+          >
+            <svg
+              className="relative hidden size-5 self-center stroke-current group-data-[selected]/option:inline sm:size-4"
+              viewBox="0 0 16 16"
+              fill="none"
+              aria-hidden="true"
+            >
+              <path
+                d="M4 8.5l3 3L12 4"
+                strokeWidth={1.5}
+                strokeLinecap="round"
+                strokeLinejoin="round"
+              />
+            </svg>
+            <span className={clsx(className, sharedClasses, 'col-start-2')}>
+              {children}
+            </span>
+          </div>
+        );
+      }}
+    </Headless.ListboxOption>
+  );
+}
+
+export function ListboxLabel({
+  className,
+  ...props
+}: React.ComponentPropsWithoutRef<'span'>) {
+  return (
+    <span
+      {...props}
+      className={clsx(
+        className,
+        'ml-2.5 truncate first:ml-0 sm:ml-2 sm:first:ml-0'
+      )}
+    />
+  );
+}
+
+export function ListboxDescription({
+  className,
+  children,
+  ...props
+}: React.ComponentPropsWithoutRef<'span'>) {
+  return (
+    <span
+      {...props}
+      className={clsx(
+        className,
+        'flex flex-1 overflow-hidden text-zinc-500 before:w-2 before:min-w-0 before:shrink group-data-[focus]/option:text-white dark:text-zinc-400'
+      )}
+    >
+      <span className="flex-1 truncate">{children}</span>
+    </span>
+  );
+}

app/components/pipeline/GeneratePipeline.tsx

@@ -13,6 +13,7 @@ import { Field, FieldGroup, Label } from '~/components/catalyst/fieldset';
 import { Textarea } from '~/components/catalyst/textarea';
 import { Spinner } from '~/components/Spinner';
 import { useLitlytics } from '~/store/WithLitLytics';
+import { Listbox, ListboxOption } from '../catalyst/listbox';
 import { RefinePipeline } from './RefinePipeline';
 
 const tabClass = clsx(
@@ -26,6 +27,7 @@ const tabClass = clsx(
 export default function GeneratePipeline() {
   const { litlytics, pipeline, setPipeline, pipelineStatus } = useLitlytics();
   const [selectedTab, setSelectedTab] = useState(0);
+  const [selectedTask, setSelectedTask] = useState<string>('');
   const [isOpen, setIsOpen] = useState(false);
   const [loading, setLoading] = useState(false);
   const [error, setError] = useState<Error>();
@@ -46,6 +48,23 @@ export default function GeneratePipeline() {
     }
   };
 
+  const suggestTasks = async () => {
+    try {
+      setLoading(true);
+      setError(undefined);
+
+      // generate plan from LLM
+      const newPipeline = await litlytics.suggestTasks();
+      setPipeline(newPipeline);
+      setSelectedTask(newPipeline.pipelineTasks?.at(0) ?? '');
+      setLoading(false);
+      setSelectedTab(1);
+    } catch (err) {
+      setError(err as Error);
+      setLoading(false);
+    }
+  };
+
   const closeDialog = () => {
     if (
       loading ||
@@ -108,16 +127,52 @@ export default function GeneratePipeline() {
                     />
                   </Field>
                 </FieldGroup>
-                <div className="flex justify-end">
-                  <Button onClick={runPlan} disabled={loading} className="mt-2">
-                    {loading && (
-                      <div className="flex items-center">
-                        <Spinner className="h-5 w-5" />
-                      </div>
-                    )}
+                <div className="flex justify-between items-center mt-2">
+                  {!pipeline.pipelineTasks?.length && (
+                    <Button onClick={suggestTasks} plain disabled={loading}>
+                      Suggest tasks using test docs
+                    </Button>
+                  )}
+                  {loading && (
+                    <div className="flex items-center">
+                      <Spinner className="h-5 w-5" />
+                    </div>
+                  )}
+                  <Button onClick={runPlan} disabled={loading}>
                     Plan
                   </Button>
                 </div>
+                {pipeline.pipelineTasks?.length && (
+                  <div className="mt-4">
+                    <Field>
+                      <Label>Suggested tasks</Label>
+                      <Listbox
+                        className="!mt-1"
+                        optionsClassName="z-50"
+                        value={selectedTask}
+                        onChange={(e) => setSelectedTask(e)}
+                      >
+                        {pipeline.pipelineTasks.map((task) => (
+                          <ListboxOption key={task} value={task}>
+                            {task}
+                          </ListboxOption>
+                        ))}
+                      </Listbox>
+                    </Field>
+
+                    <Button
+                      onClick={() =>
+                        setPipeline({
+                          ...pipeline,
+                          pipelineDescription: selectedTask,
+                        })
+                      }
+                      disabled={loading}
+                    >
+                      Use suggested task
+                    </Button>
+                  </div>
+                )}
                 {error && (
                   <div className="flex items-center justify-between bg-red-400 dark:bg-red-700 rounded-xl py-1 px-2 my-2">
                     Error planning: {error.message}

bun.lockb

@@ -6,4 +6,5 @@ export interface Doc {
   content: string;
   test?: boolean;
   processingResults: StepResult[];
+  summary?: string;
 }

packages/litlytics/doc/Document.ts

@@ -22,6 +22,7 @@ import {
   type PipelineStatus,
 } from './pipeline/Pipeline';
 import { refinePipeline } from './pipeline/refine';
+import { suggestTasks } from './pipeline/suggestTasks';
 import { generateCodeExplain } from './step/explain';
 import { generateStep, type GenerateStepArgs } from './step/generate';
 import { refineStep } from './step/refine';
@@ -251,6 +252,35 @@ export class LitLytics {
     });
   };
 
+  suggestTasks = async () => {
+    if (!this.pipeline.source?.docs?.length) {
+      return this.pipeline;
+    }
+
+    const res = await suggestTasks({
+      litlytics: this,
+      pipeline: this.pipeline,
+    });
+
+    const newDocs = this.pipeline.source.docs.map((d) => {
+      // try to find updated doc
+      const upDoc = res.docs.find((newD) => newD.id === d.id);
+      // if exists - return it
+      if (upDoc) {
+        return upDoc;
+      }
+      // otherwise - use original
+      return d;
+    });
+    // update docs
+    this.setDocs(newDocs);
+
+    // update pipeline tasks
+    return this.setPipeline({
+      pipelineTasks: res.tasks,
+    });
+  };
+
   refinePipeline = async ({ refineRequest }: { refineRequest: string }) => {
     const plan = await refinePipeline({
       litlytics: this,

packages/litlytics/litlytics.ts

@@ -80,6 +80,8 @@ export async function executeOnLLM({
   // set key
   const modelObj = getModel({ provider, model, key });
   const { text, usage } = await generateText({
+    // default to 4k max tokens output which is what most models support
+    maxTokens: 4096,
     ...modelArgs,
     model: modelObj,
     messages,

packages/litlytics/llm/llm.ts

@@ -6,6 +6,7 @@ export interface Pipeline {
   name: string;
   pipelineDescription?: string;
   pipelinePlan?: string;
+  pipelineTasks?: string[];
   source: SourceStep;
   resultDocs?: Doc[];
   results?: Result[];

packages/litlytics/pipeline/Pipeline.ts

@@ -0,0 +1,7 @@
+export const docToDescriptionPrompt = `You are an expert at analyzing documents.
+
+Your task is to analyze given document and describe it succinctly.
+The information about the document will be used to suggest data science tasks that can be done with it, so make sure to include all relevant information in your response.
+Do not include list of tasks into your response.
+
+Think the request through step-by-step inside <thinking> tags, and then provide your final response inside <output> tags.`;