Start implementing new agentic layout

Implement basic chat and first agent function that suggests pipelines
based on user documents

Addresses #21

app/components/Background.tsx

@@ -2,7 +2,8 @@ import clsx from 'clsx';
 
 export function Background({
   children,
-  className,
+  // size
+  className = 'w-screen h-screen p-6 overflow-auto',
 }: {
   children: React.ReactNode;
   className?: string;
@@ -11,8 +12,6 @@ export function Background({
     <div
       className={clsx(
         className,
-        // size
-        'w-screen h-screen p-6 overflow-auto',
         // content positioning
         'flex flex-col items-center',
         // bg dots

app/components/agent/Agent.tsx

@@ -0,0 +1,15 @@
+import { PipelineBuilder } from '../pipeline/PipelineBuilder';
+import { Chat } from './Chat';
+
+export function AgentUI() {
+  return (
+    <div className="flex flex-1 w-screen h-screen">
+      <div className="flex flex-1">
+        <Chat />
+      </div>
+      <div className="flex flex-1">
+        <PipelineBuilder className="w-full h-full pt-6 p-3" />
+      </div>
+    </div>
+  );
+}

app/components/agent/Chat.tsx

@@ -0,0 +1,103 @@
+import { PaperAirplaneIcon } from '@heroicons/react/24/solid';
+import { useAtomValue, useSetAtom } from 'jotai';
+import { useState } from 'react';
+import { litlyticsAtom, pipelineAtom } from '~/store/store';
+import { Button } from '../catalyst/button';
+import { Input } from '../catalyst/input';
+import { CustomMarkdown } from '../markdown/Markdown';
+import { askAgent, Message } from './agent';
+
+function MessageRender({ message }: { message: Message }) {
+  if (message.from === 'user') {
+    return (
+      <div className="bg-neutral-100 dark:bg-neutral-900 p-2 rounded-xl w-fit self-end">
+        {message.text}
+      </div>
+    );
+  }
+
+  return (
+    <div className="flex gap-3">
+      <div className="w-fit">
+        <span className="rounded-full border p-1 border-neutral-300 dark:border-neutral-700">
+          🔥
+        </span>
+      </div>
+      <div className="flex flex-1">
+        <div className="prose dark:prose-invert">
+          <CustomMarkdown>{message.text}</CustomMarkdown>
+        </div>
+      </div>
+    </div>
+  );
+}
+
+export function Chat() {
+  const litlytics = useAtomValue(litlyticsAtom);
+  const setPipeline = useSetAtom(pipelineAtom);
+  const [input, setInput] = useState<string>('');
+  const [messages, setMessages] = useState<Message[]>([
+    {
+      id: '0',
+      from: 'assistant',
+      text: `Hi! I'm Lit. Ask me to do anything for you.`,
+    },
+  ]);
+
+  const sendMessage = async () => {
+    const inputMessage = input;
+    // reset input
+    setInput('');
+    // append user message to messages
+    const messagesWithUser: Message[] = [
+      ...messages,
+      {
+        id: String(messages.length),
+        from: 'user',
+        text: inputMessage,
+      },
+    ];
+    setMessages(messagesWithUser);
+
+    // TODO: show loading state
+
+    // run new messages through agent
+    const newMessages = await askAgent({
+      messages: messagesWithUser,
+      litlytics,
+      setPipeline,
+    });
+    setMessages(newMessages);
+  };
+
+  return (
+    <div className="flex flex-col w-full h-full">
+      <div className="flex flex-1 flex-col gap-4 p-3 pt-20">
+        {messages.map((m) => (
+          <MessageRender key={m.id} message={m} />
+        ))}
+      </div>
+      <div className="flex items-center min-h-16 p-2">
+        <Input
+          wrapperClassName="h-fit after:hidden sm:after:focus-within:ring-2 sm:after:focus-within:ring-blue-500"
+          className="rounded-r-none"
+          placeholder="Ask Lit to do things for you"
+          value={input}
+          onChange={(e) => setInput(e.target.value)}
+          onKeyDown={(e) => {
+            if (e.key === 'Enter') {
+              sendMessage();
+            }
+          }}
+        />
+        <Button
+          className="h-9 rounded-l-none"
+          title="Send"
+          onClick={sendMessage}
+        >
+          <PaperAirplaneIcon />
+        </Button>
+      </div>
+    </div>
+  );
+}

app/components/agent/agent.ts

@@ -0,0 +1,112 @@
+import { Pipeline, tool, type LLMArgs, type LitLytics } from 'litlytics';
+import { RunPromptFromMessagesArgs } from 'litlytics/engine/runPrompt';
+import { z } from 'zod';
+
+const systemPrompt = `
+You are Lit - a friendly assistant and an expert in data science.
+
+Your task is to help user design a text document processing pipeline using low-code platform called LitLytics.
+LitLytics allows creating custom text document processing pipelines using custom processing steps.
+
+You have access to following LitLytics functions:
+- Suggest a list of possible pipelines that can be applied to user's documents
+- Generate a suggested pipeline for processing documents
+- Refine suggested pipeline for processing documents
+- Add a new step to pipeline
+- Edit a step in the pipeline
+- Test a step in the pipeline
+- Execute a pipeline
+
+If you can execute one of the functions listed above - do so and let user know you are on it.
+`;
+
+export interface Message {
+  id: string;
+  from: 'user' | 'assistant';
+  text: string;
+}
+
+export const askAgent = async ({
+  messages,
+  litlytics,
+  setPipeline,
+}: {
+  messages: Message[];
+  litlytics: LitLytics;
+  setPipeline: (p: Pipeline) => void;
+}): Promise<Message[]> => {
+  // create a promise we will use as result
+  const { promise, resolve, reject } = Promise.withResolvers<Message[]>();
+
+  // generate input messages
+  const inputMessages: RunPromptFromMessagesArgs['messages'] = messages.map(
+    (m) => ({
+      content: m.text,
+      role: m.from,
+    })
+  );
+  const agentMessages: RunPromptFromMessagesArgs['messages'] = [
+    {
+      role: 'system',
+      content: systemPrompt.trim(),
+    },
+    ...inputMessages,
+  ];
+  console.log(agentMessages);
+
+  // generate tools
+  const tools: LLMArgs['tools'] = {
+    analyzeDocuments: tool({
+      description: `Suggest a list of possible pipelines that can be applied to user's documents.`,
+      parameters: z.object({
+        suggest: z.boolean(),
+      }),
+      execute: async () => {
+        // run task
+        const newPipeline = await litlytics.suggestTasks();
+        setPipeline(newPipeline);
+        // generate a response
+        const agentMessagesWithResult = agentMessages.concat([
+          {
+            content: `Suggested tasks from function execution:
+${newPipeline.pipelineTasks?.map((task) => '- ' + task)?.join('\n')}
+
+Generate a text description for user.`,
+            role: 'system',
+          },
+        ]);
+        const result = await litlytics.runPromptFromMessages({
+          messages: agentMessagesWithResult,
+        });
+        resolve(
+          messages.concat({
+            id: String(messages.length),
+            from: 'assistant',
+            text: result.result,
+          })
+        );
+      },
+    }),
+  };
+
+  // execute request
+  const result = await litlytics.runPromptFromMessages({
+    messages: agentMessages,
+    args: {
+      tools,
+    },
+  });
+
+  console.log(result);
+  if (result.result.length) {
+    resolve(
+      messages.concat({
+        id: String(messages.length),
+        from: 'assistant',
+        text: result.result,
+      })
+    );
+  }
+
+  return promise;
+};

app/components/pipeline/PipelineBuilder.tsx

@@ -4,11 +4,11 @@ import { OutputNode } from './nodes/output/OutputNode';
 import { SourceNode } from './nodes/source/SourceNode';
 import { StepNode } from './nodes/StepNode';
 
-export function PipelineBuilder() {
+export function PipelineBuilder({ className }: { className?: string }) {
   const { pipeline } = useLitlytics();
 
   return (
-    <Background>
+    <Background className={className}>
       <SourceNode />
       {pipeline.steps
         .sort((a, b) => (a.connectsTo.includes(b.id) ? -1 : 1))

app/components/ui/Overlay.tsx

@@ -3,9 +3,11 @@ import {
   ArrowUturnLeftIcon,
   ArrowUturnRightIcon,
   Bars3Icon,
+  ChatBubbleBottomCenterIcon,
   Cog8ToothIcon,
   FolderIcon,
   QuestionMarkCircleIcon,
+  QueueListIcon,
   TrashIcon,
 } from '@heroicons/react/24/solid';
 import { useAtom, useAtomValue } from 'jotai';
@@ -29,7 +31,7 @@ import {
   DropdownMenu,
   DropdownShortcut,
 } from '~/components/catalyst/dropdown';
-import { configAtom, pipelineUndoAtom } from '~/store/store';
+import { configAtom, pipelineUndoAtom, uiLayoutAtom } from '~/store/store';
 import { useLitlytics } from '~/store/WithLitLytics';
 import { Field, FieldGroup, Label } from '../catalyst/fieldset';
 import { Input } from '../catalyst/input';
@@ -66,6 +68,7 @@ export function OverlayUI() {
     useLitlytics();
   const litlyticsConfig = useAtomValue(configAtom);
   const { undo, redo, canUndo, canRedo } = useAtomValue(pipelineUndoAtom);
+  const [uiLayout, setUiLayout] = useAtom(uiLayoutAtom);
   const [isOpen, setIsOpen] = useState(false);
   const [isSaveOpen, setIsSaveOpen] = useState(false);
   const [isLoadOpen, setIsLoadOpen] = useState(false);
@@ -266,6 +269,21 @@ export function OverlayUI() {
             </DropdownMenu>
           </Dropdown>
         </MenuHolder>
+
+        <div className="pointer-events-auto">
+          <Button
+            outline
+            onClick={() =>
+              setUiLayout((l) => (l === 'agent' ? 'execution' : 'agent'))
+            }
+          >
+            {uiLayout === 'agent' ? (
+              <QueueListIcon title="Switch to execution UI" />
+            ) : (
+              <ChatBubbleBottomCenterIcon title="Switch to agent UI" />
+            )}
+          </Button>
+        </div>
       </div>
 
       {/* Pipeline reset dialog */}

app/components/ui/UI.tsx

@@ -1,7 +1,10 @@
+import { useAtomValue } from 'jotai';
 import { useEffect, useState } from 'react';
 import { PipelineBuilder } from '~/components/pipeline/PipelineBuilder';
 import { OverlayUI } from '~/components/ui/Overlay';
+import { uiLayoutAtom } from '~/store/store';
 import { WithLitLytics } from '~/store/WithLitLytics';
+import { AgentUI } from '../agent/Agent';
 import { Background } from '../Background';
 import { Spinner } from '../Spinner';
 
@@ -33,10 +36,13 @@ function ClientOnly({ children }: { children: React.ReactNode }) {
 }
 
 export function UI() {
+  const uiLayout = useAtomValue(uiLayoutAtom);
+
   return (
     <ClientOnly>
       <OverlayUI />
-      <PipelineBuilder />
+      {uiLayout === 'execution' && <PipelineBuilder />}
+      {uiLayout === 'agent' && <AgentUI />}
     </ClientOnly>
   );
 }

app/store/store.ts

@@ -42,3 +42,11 @@ export const litlyticsAtom = atom<LitLytics>(
     key: '',
   })
 );
+
+// litlytics config, persisted in localStorage
+export const uiLayoutAtom = atomWithStorage<'agent' | 'execution'>(
+  'litlytics.uilayout',
+  'agent',
+  undefined,
+  { getOnInit: true }
+);

package.json

@@ -38,7 +38,8 @@
     "rehype-raw": "^7.0.0",
     "remark-breaks": "^4.0.0",
     "remark-gfm": "^4.0.0",
-    "turndown": "^7.2.0"
+    "turndown": "^7.2.0",
+    "zod": "^3.23.8"
   },
   "devDependencies": {
     "@remix-run/dev": "^2.13.1",

packages/litlytics/litlytics.ts

@@ -33,6 +33,7 @@ import {
 } from './step/Step';
 
 // export types and commonly used vars
+export { tool } from 'ai';
 export type { Doc } from './doc/Document';
 export { modelCosts } from './llm/costs';
 export {
@@ -268,6 +269,11 @@ export class LitLytics {
       return this.pipeline;
     }
 
+    // do not run suggestion if the tasks are already generated
+    if (this.pipeline.pipelineTasks?.length) {
+      return this.pipeline;
+    }
+
     const res = await suggestTasks({
       litlytics: this,
       pipeline: this.pipeline,