Add new agent / chat UI

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 overflow-auto pt-6 p-3" />
+      </div>
+    </div>
+  );
+}

app/components/agent/Chat.tsx

@@ -0,0 +1,144 @@
+import { PaperAirplaneIcon } from '@heroicons/react/24/solid';
+import { useAtomValue, useSetAtom } from 'jotai';
+import { useEffect, useRef, 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 { Spinner } from '../Spinner';
+import { askAgent } from './logic/askAgent';
+import { type Message } from './logic/types';
+
+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 messageBoxRef = useRef<HTMLDivElement>(null);
+  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 [loading, setLoading] = useState<boolean>(false);
+  const [error, setError] = useState<Error | undefined>();
+
+  useEffect(() => {
+    // scroll to bottom
+    if (messageBoxRef.current) {
+      messageBoxRef.current.scrollTo({
+        top: messageBoxRef.current.scrollHeight,
+        behavior: 'smooth',
+      });
+    }
+  }, [messages]);
+
+  const sendMessage = async () => {
+    const inputMessage = input.trim();
+    // do nothing if there's no user message
+    if (!inputMessage.length) {
+      return;
+    }
+    // reset input
+    setInput('');
+    // reset error
+    setError(undefined);
+    // append user message to messages
+    const messagesWithUser: Message[] = [
+      ...messages,
+      {
+        id: String(messages.length),
+        from: 'user',
+        text: inputMessage,
+      },
+    ];
+    setMessages(messagesWithUser);
+
+    // show loading state
+    setLoading(true);
+    // run new messages through agent
+    try {
+      const newMessages = await askAgent({
+        messages: messagesWithUser,
+        litlytics,
+        setPipeline,
+      });
+      setMessages(newMessages);
+    } catch (err) {
+      // catch and display error
+      setError(err as Error);
+    }
+    // disable loading state
+    setLoading(false);
+  };
+
+  return (
+    <div className="flex flex-col w-full h-full">
+      <div
+        ref={messageBoxRef}
+        className="flex flex-1 flex-col gap-4 p-3 pt-20 max-h-screen overflow-auto"
+      >
+        {messages.map((m) => (
+          <MessageRender key={m.id} message={m} />
+        ))}
+        {loading && (
+          <div className="flex items-center justify-end gap-2">
+            <Spinner className="h-5 w-5" /> Thinking...
+          </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 while thinking: {error.message}
+          </div>
+        )}
+      </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/logic/askAgent.ts

@@ -0,0 +1,81 @@
+import { Pipeline, type LLMArgs, type LitLytics } from 'litlytics';
+import { RunPromptFromMessagesArgs } from 'litlytics/engine/runPrompt';
+import { agentSystemPrompt } from './prompts/system';
+import { agentTools } from './tools/tools';
+import { type Message } from './types';
+
+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,
+    })
+  );
+  // generate functions list
+  const functionsList = agentTools.map((t) => `- ${t.description}`).join('\n');
+  // prepend system message
+  const agentMessages: RunPromptFromMessagesArgs['messages'] = [
+    // system prompt
+    {
+      role: 'system',
+      content: agentSystemPrompt.trim().replace('{{FUNCTIONS}}', functionsList),
+    },
+    // current pipeline for context
+    {
+      role: 'system',
+      content: `Current pipeline:
+\`\`\`
+${JSON.stringify(litlytics.pipeline, null, 2)}
+\`\`\``,
+    },
+    // user messages
+    ...inputMessages,
+  ];
+  console.log(agentMessages);
+
+  // generate tools
+  const tools: LLMArgs['tools'] = {};
+  for (const tool of agentTools) {
+    tools[tool.name] = tool.create({
+      litlytics,
+      setPipeline,
+      agentMessages,
+      messages,
+      resolve,
+      reject,
+    });
+  }
+
+  // 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/agent/logic/prompts/system.ts

@@ -0,0 +1,12 @@
+export const agentSystemPrompt = `
+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.
+LitLytics supports text documents and .csv, .doc(x), .pdf, .txt text files.
+
+You have access to following LitLytics functions:
+{{FUNCTIONS}}
+
+If you can execute one of the functions listed above - do so and let user know you are on it.
+`;

app/components/agent/logic/tools/addNewStep.ts

@@ -0,0 +1,85 @@
+import { SourceStep } from '@/packages/litlytics/litlytics';
+import { ProcessingStep, tool } from 'litlytics';
+import { z } from 'zod';
+import { ToolDefinition } from '../types';
+
+const description = `Function description: Add a new step to the pipeline
+  Function arguments: step type, name, description, input type and a step to connect to
+  Extra instructions: User must specify arguments themselves. Consider primary source to be a possible source step as well.`;
+
+export const addNewStep: ToolDefinition = {
+  name: 'addNewStep',
+  description,
+  create: ({
+    litlytics,
+    setPipeline,
+    agentMessages,
+    messages,
+    resolve,
+    reject,
+  }) =>
+    tool({
+      description,
+      parameters: z.object({
+        stepType: z.enum(['llm', 'code']),
+        stepName: z.string(),
+        stepDescription: z.string(),
+        stepInput: z.enum(['doc', 'result', 'aggregate-docs', 'aggregate-results']),
+        sourceStepId: z.string().optional(),
+      }),
+      execute: async ({ stepType, stepName, stepDescription, stepInput, sourceStepId }) => {
+        try {
+          const newStep = {
+            id: crypto.randomUUID(), // Generate a unique ID for the step using UUID
+            name: stepName,
+            description: stepDescription,
+            type: stepType,
+            connectsTo: [],
+            input: stepInput,
+          };
+
+          // find source step by ID
+          let sourceStep: SourceStep | ProcessingStep | undefined = litlytics.pipeline.steps.find((s) => s.id === sourceStepId);
+          if (sourceStepId === litlytics.pipeline.source.id) {
+            sourceStep = litlytics.pipeline.source;
+          }
+
+          // add the new step to the pipeline
+          const newPipeline = await litlytics.addStep({
+            step: newStep,
+            sourceStep,
+          });
+
+          setPipeline(newPipeline);
+
+          // find newly added step
+          const createdStep = newPipeline.steps.find((s) => s.name === newStep.name);
+
+          // add a message to the agent messages
+          const agentMessagesWithResult = agentMessages.concat([
+            {
+              content: `New step added: \`\`\`
+${JSON.stringify(createdStep, null, 2)}
+\`\`\``,
+              role: 'system',
+            },
+          ]);
+
+          const result = await litlytics.runPromptFromMessages({
+            messages: agentMessagesWithResult,
+          });
+
+          resolve(
+            messages.concat({
+              id: String(messages.length),
+              from: 'assistant',
+              text: result.result,
+            })
+          );
+        } catch (err) {
+          reject(err as Error);
+        }
+      },
+    }),
+};
+