Skip to content

Commit 1255a79

Browse files
committed
Add a new agent tool to add new steps to pipeline
Addresses #21
1 parent 3d46843 commit 1255a79

File tree

7 files changed

+101
-4
lines changed

7 files changed

+101
-4
lines changed

app/components/agent/logic/askAgent.ts

+10
Original file line numberDiff line numberDiff line change
@@ -27,10 +27,20 @@ export const askAgent = async ({
2727
const functionsList = agentTools.map((t) => `- ${t.description}`).join('\n');
2828
// prepend system message
2929
const agentMessages: RunPromptFromMessagesArgs['messages'] = [
30+
// system prompt
3031
{
3132
role: 'system',
3233
content: agentSystemPrompt.trim().replace('{{FUNCTIONS}}', functionsList),
3334
},
35+
// current pipeline for context
36+
{
37+
role: 'system',
38+
content: `Current pipeline:
39+
\`\`\`
40+
${JSON.stringify(litlytics.pipeline, null, 2)}
41+
\`\`\``,
42+
},
43+
// user messages
3444
...inputMessages,
3545
];
3646
console.log(agentMessages);
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
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+

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import { tool } from 'litlytics';
22
import { z } from 'zod';
33
import { ToolDefinition } from '../types';
44

5-
const description = `Suggest a list of potential pipelines / actions that can be executed on user documents`;
5+
const description = `Function description: Suggest a list of potential pipelines / actions that can be executed on user documents`;
66

77
export const analyzeDocuments: ToolDefinition = {
88
name: 'analyzeDocuments',

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import { tool } from 'litlytics';
22
import { z } from 'zod';
33
import { ToolDefinition } from '../types';
44

5-
const description = `Assemble suggested pipeline`;
5+
const description = `Function description: Assemble suggested pipeline`;
66

77
export const createSuggestedPipeline: ToolDefinition = {
88
name: 'createSuggestedPipeline',

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import { tool } from 'litlytics';
22
import { z } from 'zod';
33
import { ToolDefinition } from '../types';
44

5-
const description = `Refine suggested pipeline using user request`;
5+
const description = `Function description: Refine suggested pipeline using user request`;
66

77
export const refinePipeline: ToolDefinition = {
88
name: 'refinePipeline',

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import { tool } from 'litlytics';
22
import { z } from 'zod';
33
import { ToolDefinition } from '../types';
44

5-
const description = `Suggest a new pipeline for processing documents for a given task from user`;
5+
const description = `Function description: Suggest a new pipeline for processing documents for a given task from user`;
66

77
export const suggestPipeline: ToolDefinition = {
88
name: 'suggestPipeline',

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

+2
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import { addNewStep } from './addNewStep';
12
import { analyzeDocuments } from './analyzeDocs';
23
import { createSuggestedPipeline } from './createSuggestedPipeline';
34
import { refinePipeline } from './refinePipeline';
@@ -8,4 +9,5 @@ export const agentTools = [
89
suggestPipeline,
910
refinePipeline,
1011
createSuggestedPipeline,
12+
addNewStep,
1113
];

0 commit comments

Comments
 (0)