Refactor agent tool code to make them modular

Assemble tools description in prompt from tools themselves as well

Author: yamalight

app/components/agent/Chat.tsx

@@ -5,7 +5,8 @@ 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';
+import { askAgent } from './logic/askAgent';
+import { type Message } from './logic/types';
 
 function MessageRender({ message }: { message: Message }) {
   if (message.from === 'user') {

app/components/agent/agent.ts

@@ -0,0 +1,71 @@
+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'] = [
+    {
+      role: 'system',
+      content: agentSystemPrompt.trim().replace('{{FUNCTIONS}}', functionsList),
+    },
+    ...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/askAgent.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/prompts/system.ts

@@ -0,0 +1,53 @@
+import { tool } from 'litlytics';
+import { z } from 'zod';
+import { ToolDefinition } from '../types';
+
+const description = `Suggest a list of potential pipelines / actions that can be executed on user documents`;
+
+export const analyzeDocuments: ToolDefinition = {
+  name: 'analyzeDocuments',
+  description,
+  create: ({
+    litlytics,
+    setPipeline,
+    agentMessages,
+    messages,
+    resolve,
+    reject,
+  }) =>
+    tool({
+      description,
+      parameters: z.object({
+        suggest: z.boolean(),
+      }),
+      execute: async () => {
+        try {
+          // 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,
+            })
+          );
+        } catch (err) {
+          reject(err as Error);
+        }
+      },
+    }),
+};