|
| 1 | +import { SourceStep } from '@/packages/litlytics/litlytics'; |
| 2 | +import { ProcessingStep, tool } from 'litlytics'; |
| 3 | +import { z } from 'zod'; |
| 4 | +import { ToolDefinition } from '../types'; |
| 5 | + |
| 6 | +const description = `Function description: Add a new step to the pipeline |
| 7 | + Function arguments: step type, name, description, input type and a step to connect to |
| 8 | + Extra instructions: User must specify arguments themselves. Consider primary source to be a possible source step as well.`; |
| 9 | + |
| 10 | +export const addNewStep: ToolDefinition = { |
| 11 | + name: 'addNewStep', |
| 12 | + description, |
| 13 | + create: ({ |
| 14 | + litlytics, |
| 15 | + setPipeline, |
| 16 | + agentMessages, |
| 17 | + messages, |
| 18 | + resolve, |
| 19 | + reject, |
| 20 | + }) => |
| 21 | + tool({ |
| 22 | + description, |
| 23 | + parameters: z.object({ |
| 24 | + stepType: z.enum(['llm', 'code']), |
| 25 | + stepName: z.string(), |
| 26 | + stepDescription: z.string(), |
| 27 | + stepInput: z.enum(['doc', 'result', 'aggregate-docs', 'aggregate-results']), |
| 28 | + sourceStepId: z.string().optional(), |
| 29 | + }), |
| 30 | + execute: async ({ stepType, stepName, stepDescription, stepInput, sourceStepId }) => { |
| 31 | + try { |
| 32 | + const newStep = { |
| 33 | + id: crypto.randomUUID(), // Generate a unique ID for the step using UUID |
| 34 | + name: stepName, |
| 35 | + description: stepDescription, |
| 36 | + type: stepType, |
| 37 | + connectsTo: [], |
| 38 | + input: stepInput, |
| 39 | + }; |
| 40 | + |
| 41 | + // find source step by ID |
| 42 | + let sourceStep: SourceStep | ProcessingStep | undefined = litlytics.pipeline.steps.find((s) => s.id === sourceStepId); |
| 43 | + if (sourceStepId === litlytics.pipeline.source.id) { |
| 44 | + sourceStep = litlytics.pipeline.source; |
| 45 | + } |
| 46 | + |
| 47 | + // add the new step to the pipeline |
| 48 | + const newPipeline = await litlytics.addStep({ |
| 49 | + step: newStep, |
| 50 | + sourceStep, |
| 51 | + }); |
| 52 | + |
| 53 | + setPipeline(newPipeline); |
| 54 | + |
| 55 | + // find newly added step |
| 56 | + const createdStep = newPipeline.steps.find((s) => s.name === newStep.name); |
| 57 | + |
| 58 | + // add a message to the agent messages |
| 59 | + const agentMessagesWithResult = agentMessages.concat([ |
| 60 | + { |
| 61 | + content: `New step added: \`\`\` |
| 62 | +${JSON.stringify(createdStep, null, 2)} |
| 63 | +\`\`\``, |
| 64 | + role: 'system', |
| 65 | + }, |
| 66 | + ]); |
| 67 | + |
| 68 | + const result = await litlytics.runPromptFromMessages({ |
| 69 | + messages: agentMessagesWithResult, |
| 70 | + }); |
| 71 | + |
| 72 | + resolve( |
| 73 | + messages.concat({ |
| 74 | + id: String(messages.length), |
| 75 | + from: 'assistant', |
| 76 | + text: result.result, |
| 77 | + }) |
| 78 | + ); |
| 79 | + } catch (err) { |
| 80 | + reject(err as Error); |
| 81 | + } |
| 82 | + }, |
| 83 | + }), |
| 84 | +}; |
| 85 | + |
0 commit comments