Skip to content

Commit 35b2266

Browse files
committed
Refactor agent tool code to make them modular
Assemble tools description in prompt from tools themselves as well
1 parent 8cd9738 commit 35b2266

File tree

12 files changed

+349
-221
lines changed

12 files changed

+349
-221
lines changed

app/components/agent/Chat.tsx

+2-1
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,8 @@ import { litlyticsAtom, pipelineAtom } from '~/store/store';
55
import { Button } from '../catalyst/button';
66
import { Input } from '../catalyst/input';
77
import { CustomMarkdown } from '../markdown/Markdown';
8-
import { askAgent, Message } from './agent';
8+
import { askAgent } from './logic/askAgent';
9+
import { type Message } from './logic/types';
910

1011
function MessageRender({ message }: { message: Message }) {
1112
if (message.from === 'user') {

app/components/agent/agent.ts

-219
This file was deleted.
+71
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
import { Pipeline, type LLMArgs, type LitLytics } from 'litlytics';
2+
import { RunPromptFromMessagesArgs } from 'litlytics/engine/runPrompt';
3+
import { agentSystemPrompt } from './prompts/system';
4+
import { agentTools } from './tools/tools';
5+
import { type Message } from './types';
6+
7+
export const askAgent = async ({
8+
messages,
9+
litlytics,
10+
setPipeline,
11+
}: {
12+
messages: Message[];
13+
litlytics: LitLytics;
14+
setPipeline: (p: Pipeline) => void;
15+
}): Promise<Message[]> => {
16+
// create a promise we will use as result
17+
const { promise, resolve, reject } = Promise.withResolvers<Message[]>();
18+
19+
// generate input messages
20+
const inputMessages: RunPromptFromMessagesArgs['messages'] = messages.map(
21+
(m) => ({
22+
content: m.text,
23+
role: m.from,
24+
})
25+
);
26+
// generate functions list
27+
const functionsList = agentTools.map((t) => `- ${t.description}`).join('\n');
28+
// prepend system message
29+
const agentMessages: RunPromptFromMessagesArgs['messages'] = [
30+
{
31+
role: 'system',
32+
content: agentSystemPrompt.trim().replace('{{FUNCTIONS}}', functionsList),
33+
},
34+
...inputMessages,
35+
];
36+
console.log(agentMessages);
37+
38+
// generate tools
39+
const tools: LLMArgs['tools'] = {};
40+
for (const tool of agentTools) {
41+
tools[tool.name] = tool.create({
42+
litlytics,
43+
setPipeline,
44+
agentMessages,
45+
messages,
46+
resolve,
47+
reject,
48+
});
49+
}
50+
51+
// execute request
52+
const result = await litlytics.runPromptFromMessages({
53+
messages: agentMessages,
54+
args: {
55+
tools,
56+
},
57+
});
58+
59+
console.log(result);
60+
if (result.result.length) {
61+
resolve(
62+
messages.concat({
63+
id: String(messages.length),
64+
from: 'assistant',
65+
text: result.result,
66+
})
67+
);
68+
}
69+
70+
return promise;
71+
};
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
export const agentSystemPrompt = `
2+
You are Lit - a friendly assistant and an expert in data science.
3+
4+
Your task is to help user design a text document processing pipeline using low-code platform called LitLytics.
5+
LitLytics allows creating custom text document processing pipelines using custom processing steps.
6+
LitLytics supports text documents and .csv, .doc(x), .pdf, .txt text files.
7+
8+
You have access to following LitLytics functions:
9+
{{FUNCTIONS}}
10+
11+
If you can execute one of the functions listed above - do so and let user know you are on it.
12+
`;
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
import { tool } from 'litlytics';
2+
import { z } from 'zod';
3+
import { ToolDefinition } from '../types';
4+
5+
const description = `Suggest a list of potential pipelines / actions that can be executed on user documents`;
6+
7+
export const analyzeDocuments: ToolDefinition = {
8+
name: 'analyzeDocuments',
9+
description,
10+
create: ({
11+
litlytics,
12+
setPipeline,
13+
agentMessages,
14+
messages,
15+
resolve,
16+
reject,
17+
}) =>
18+
tool({
19+
description,
20+
parameters: z.object({
21+
suggest: z.boolean(),
22+
}),
23+
execute: async () => {
24+
try {
25+
// run task
26+
const newPipeline = await litlytics.suggestTasks();
27+
setPipeline(newPipeline);
28+
// generate a response
29+
const agentMessagesWithResult = agentMessages.concat([
30+
{
31+
content: `Suggested tasks from function execution:
32+
${newPipeline.pipelineTasks?.map((task) => '- ' + task)?.join('\n')}
33+
34+
Generate a text description for user.`,
35+
role: 'system',
36+
},
37+
]);
38+
const result = await litlytics.runPromptFromMessages({
39+
messages: agentMessagesWithResult,
40+
});
41+
resolve(
42+
messages.concat({
43+
id: String(messages.length),
44+
from: 'assistant',
45+
text: result.result,
46+
})
47+
);
48+
} catch (err) {
49+
reject(err as Error);
50+
}
51+
},
52+
}),
53+
};

0 commit comments

Comments
 (0)