Skip to content

Commit 4b87896

Browse files
yamalightyamalight
authored
Merge pull request #22 from yamalight/feature/agent
Add new agent / chat UI
2 parents 58de1c7 + 391dc38 commit 4b87896

22 files changed

+879
-11
lines changed

‎app/components/Background.tsx

+2-3
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,8 @@ import clsx from 'clsx';
22

33
export function Background({
44
children,
5-
className,
5+
// size
6+
className = 'w-screen h-screen p-6 overflow-auto',
67
}: {
78
children: React.ReactNode;
89
className?: string;
@@ -11,8 +12,6 @@ export function Background({
1112
<div
1213
className={clsx(
1314
className,
14-
// size
15-
'w-screen h-screen p-6 overflow-auto',
1615
// content positioning
1716
'flex flex-col items-center',
1817
// bg dots

‎app/components/agent/Agent.tsx

+15
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
import { PipelineBuilder } from '../pipeline/PipelineBuilder';
2+
import { Chat } from './Chat';
3+
4+
export function AgentUI() {
5+
return (
6+
<div className="flex flex-1 w-screen h-screen">
7+
<div className="flex flex-1">
8+
<Chat />
9+
</div>
10+
<div className="flex flex-1">
11+
<PipelineBuilder className="w-full h-full overflow-auto pt-6 p-3" />
12+
</div>
13+
</div>
14+
);
15+
}

‎app/components/agent/Chat.tsx

+144
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
1+
import { PaperAirplaneIcon } from '@heroicons/react/24/solid';
2+
import { useAtomValue, useSetAtom } from 'jotai';
3+
import { useEffect, useRef, useState } from 'react';
4+
import { litlyticsAtom, pipelineAtom } from '~/store/store';
5+
import { Button } from '../catalyst/button';
6+
import { Input } from '../catalyst/input';
7+
import { CustomMarkdown } from '../markdown/Markdown';
8+
import { Spinner } from '../Spinner';
9+
import { askAgent } from './logic/askAgent';
10+
import { type Message } from './logic/types';
11+
12+
function MessageRender({ message }: { message: Message }) {
13+
if (message.from === 'user') {
14+
return (
15+
<div className="bg-neutral-100 dark:bg-neutral-900 p-2 rounded-xl w-fit self-end">
16+
{message.text}
17+
</div>
18+
);
19+
}
20+
21+
return (
22+
<div className="flex gap-3">
23+
<div className="w-fit">
24+
<span className="rounded-full border p-1 border-neutral-300 dark:border-neutral-700">
25+
🔥
26+
</span>
27+
</div>
28+
<div className="flex flex-1">
29+
<div className="prose dark:prose-invert">
30+
<CustomMarkdown>{message.text}</CustomMarkdown>
31+
</div>
32+
</div>
33+
</div>
34+
);
35+
}
36+
37+
export function Chat() {
38+
const messageBoxRef = useRef<HTMLDivElement>(null);
39+
const litlytics = useAtomValue(litlyticsAtom);
40+
const setPipeline = useSetAtom(pipelineAtom);
41+
const [input, setInput] = useState<string>('');
42+
const [messages, setMessages] = useState<Message[]>([
43+
{
44+
id: '0',
45+
from: 'assistant',
46+
text: `Hi! I'm Lit. Ask me to do anything for you.`,
47+
},
48+
]);
49+
const [loading, setLoading] = useState<boolean>(false);
50+
const [error, setError] = useState<Error | undefined>();
51+
52+
useEffect(() => {
53+
// scroll to bottom
54+
if (messageBoxRef.current) {
55+
messageBoxRef.current.scrollTo({
56+
top: messageBoxRef.current.scrollHeight,
57+
behavior: 'smooth',
58+
});
59+
}
60+
}, [messages]);
61+
62+
const sendMessage = async () => {
63+
const inputMessage = input.trim();
64+
// do nothing if there's no user message
65+
if (!inputMessage.length) {
66+
return;
67+
}
68+
// reset input
69+
setInput('');
70+
// reset error
71+
setError(undefined);
72+
// append user message to messages
73+
const messagesWithUser: Message[] = [
74+
...messages,
75+
{
76+
id: String(messages.length),
77+
from: 'user',
78+
text: inputMessage,
79+
},
80+
];
81+
setMessages(messagesWithUser);
82+
83+
// show loading state
84+
setLoading(true);
85+
// run new messages through agent
86+
try {
87+
const newMessages = await askAgent({
88+
messages: messagesWithUser,
89+
litlytics,
90+
setPipeline,
91+
});
92+
setMessages(newMessages);
93+
} catch (err) {
94+
// catch and display error
95+
setError(err as Error);
96+
}
97+
// disable loading state
98+
setLoading(false);
99+
};
100+
101+
return (
102+
<div className="flex flex-col w-full h-full">
103+
<div
104+
ref={messageBoxRef}
105+
className="flex flex-1 flex-col gap-4 p-3 pt-20 max-h-screen overflow-auto"
106+
>
107+
{messages.map((m) => (
108+
<MessageRender key={m.id} message={m} />
109+
))}
110+
{loading && (
111+
<div className="flex items-center justify-end gap-2">
112+
<Spinner className="h-5 w-5" /> Thinking...
113+
</div>
114+
)}
115+
{error && (
116+
<div className="flex items-center justify-between bg-red-400 dark:bg-red-700 rounded-xl py-1 px-2 my-2">
117+
Error while thinking: {error.message}
118+
</div>
119+
)}
120+
</div>
121+
<div className="flex items-center min-h-16 p-2">
122+
<Input
123+
wrapperClassName="h-fit after:hidden sm:after:focus-within:ring-2 sm:after:focus-within:ring-blue-500"
124+
className="rounded-r-none"
125+
placeholder="Ask Lit to do things for you"
126+
value={input}
127+
onChange={(e) => setInput(e.target.value)}
128+
onKeyDown={(e) => {
129+
if (e.key === 'Enter') {
130+
sendMessage();
131+
}
132+
}}
133+
/>
134+
<Button
135+
className="h-9 rounded-l-none"
136+
title="Send"
137+
onClick={sendMessage}
138+
>
139+
<PaperAirplaneIcon />
140+
</Button>
141+
</div>
142+
</div>
143+
);
144+
}

‎app/components/agent/logic/askAgent.ts

+81
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
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+
// system prompt
31+
{
32+
role: 'system',
33+
content: agentSystemPrompt.trim().replace('{{FUNCTIONS}}', functionsList),
34+
},
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
44+
...inputMessages,
45+
];
46+
console.log(agentMessages);
47+
48+
// generate tools
49+
const tools: LLMArgs['tools'] = {};
50+
for (const tool of agentTools) {
51+
tools[tool.name] = tool.create({
52+
litlytics,
53+
setPipeline,
54+
agentMessages,
55+
messages,
56+
resolve,
57+
reject,
58+
});
59+
}
60+
61+
// execute request
62+
const result = await litlytics.runPromptFromMessages({
63+
messages: agentMessages,
64+
args: {
65+
tools,
66+
},
67+
});
68+
69+
console.log(result);
70+
if (result.result.length) {
71+
resolve(
72+
messages.concat({
73+
id: String(messages.length),
74+
from: 'assistant',
75+
text: result.result,
76+
})
77+
);
78+
}
79+
80+
return promise;
81+
};

‎app/components/agent/logic/prompts/system.ts

+12
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+
`;

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

+85
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+

0 commit comments

Comments
 (0)