Add a new agent tool to add new steps to pipeline

Addresses #21

Author: yamalight

app/components/agent/logic/askAgent.ts

@@ -27,10 +27,20 @@ export const askAgent = async ({
   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);

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);
+        }
+      },
+    }),
+};
+

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

@@ -2,7 +2,7 @@ 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`;
+const description = `Function description: Suggest a list of potential pipelines / actions that can be executed on user documents`;
 
 export const analyzeDocuments: ToolDefinition = {
   name: 'analyzeDocuments',

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

@@ -2,7 +2,7 @@ import { tool } from 'litlytics';
 import { z } from 'zod';
 import { ToolDefinition } from '../types';
 
-const description = `Assemble suggested pipeline`;
+const description = `Function description: Assemble suggested pipeline`;
 
 export const createSuggestedPipeline: ToolDefinition = {
   name: 'createSuggestedPipeline',

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

@@ -2,7 +2,7 @@ import { tool } from 'litlytics';
 import { z } from 'zod';
 import { ToolDefinition } from '../types';
 
-const description = `Refine suggested pipeline using user request`;
+const description = `Function description: Refine suggested pipeline using user request`;
 
 export const refinePipeline: ToolDefinition = {
   name: 'refinePipeline',

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

@@ -2,7 +2,7 @@ import { tool } from 'litlytics';
 import { z } from 'zod';
 import { ToolDefinition } from '../types';
 
-const description = `Suggest a new pipeline for processing documents for a given task from user`;
+const description = `Function description: Suggest a new pipeline for processing documents for a given task from user`;
 
 export const suggestPipeline: ToolDefinition = {
   name: 'suggestPipeline',

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

@@ -1,3 +1,4 @@
+import { addNewStep } from './addNewStep';
 import { analyzeDocuments } from './analyzeDocs';
 import { createSuggestedPipeline } from './createSuggestedPipeline';
 import { refinePipeline } from './refinePipeline';
@@ -8,4 +9,5 @@ export const agentTools = [
   suggestPipeline,
   refinePipeline,
   createSuggestedPipeline,
+  addNewStep,
 ];