82 lines
2.4 KiB
TypeScript
82 lines
2.4 KiB
TypeScript
|
import { WebSocket } from 'ws';
|
||
|
|
import pickSuitableAgent from '../core/agentPicker';
|
||
|
|
import handleWebSearch from '../agents/webSearchAgent';
|
||
|
|
import { BaseMessage, AIMessage, HumanMessage } from '@langchain/core/messages';
|
||
|
|
|
||
|
|
type Message = {
|
||
|
|
type: string;
|
||
|
|
content: string;
|
||
|
|
copilot: boolean;
|
||
|
|
focus: string;
|
||
|
|
history: Array<[string, string]>;
|
||
|
|
};
|
||
|
|
|
||
|
|
export const handleMessage = async (message: string, ws: WebSocket) => {
|
||
|
|
try {
|
||
|
|
const parsedMessage = JSON.parse(message) as Message;
|
||
|
|
const id = Math.random().toString(36).substring(7);
|
||
|
|
|
||
|
|
if (!parsedMessage.content)
|
||
|
|
return ws.send(
|
||
|
|
JSON.stringify({ type: 'error', data: 'Invalid message format' }),
|
||
|
|
);
|
||
|
|
|
||
|
|
const history: BaseMessage[] = parsedMessage.history.map((msg) => {
|
||
|
|
if (msg[0] === 'human') {
|
||
|
|
return new HumanMessage({
|
||
|
|
content: msg[1],
|
||
|
|
});
|
||
|
|
} else {
|
||
|
|
return new AIMessage({
|
||
|
|
content: msg[1],
|
||
|
|
});
|
||
|
|
}
|
||
|
|
});
|
||
|
|
|
||
|
|
if (parsedMessage.type === 'message') {
|
||
|
|
/* if (!parsedMessage.focus) {
|
||
|
|
const agent = await pickSuitableAgent(parsedMessage.content);
|
||
|
|
parsedMessage.focus = agent;
|
||
|
|
} */
|
||
|
|
|
||
|
|
parsedMessage.focus = 'webSearch';
|
||
|
|
|
||
|
|
switch (parsedMessage.focus) {
|
||
|
|
case 'webSearch': {
|
||
|
|
const emitter = handleWebSearch(parsedMessage.content, history);
|
||
|
|
emitter.on('data', (data) => {
|
||
|
|
const parsedData = JSON.parse(data);
|
||
|
|
if (parsedData.type === 'response') {
|
||
|
|
ws.send(
|
||
|
|
JSON.stringify({
|
||
|
|
type: 'message',
|
||
|
|
data: parsedData.data,
|
||
|
|
messageId: id,
|
||
|
|
}),
|
||
|
|
);
|
||
|
|
} else if (parsedData.type === 'sources') {
|
||
|
|
ws.send(
|
||
|
|
JSON.stringify({
|
||
|
|
type: 'sources',
|
||
|
|
data: parsedData.data,
|
||
|
|
messageId: id,
|
||
|
|
}),
|
||
|
|
);
|
||
|
|
}
|
||
|
|
});
|
||
|
|
emitter.on('end', () => {
|
||
|
|
ws.send(JSON.stringify({ type: 'messageEnd', messageId: id }));
|
||
|
|
});
|
||
|
|
emitter.on('error', (data) => {
|
||
|
|
const parsedData = JSON.parse(data);
|
||
|
|
ws.send(JSON.stringify({ type: 'error', data: parsedData.data }));
|
||
|
|
});
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
} catch (error) {
|
||
|
|
console.error('Failed to handle message', error);
|
||
|
|
ws.send(JSON.stringify({ type: 'error', data: 'Invalid message format' }));
|
||
|
|
}
|
||
|
|
};
|