Initial commit

src/websocket/messageHandler.ts

@@ -0,0 +1,81 @@
+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' }));
+  }
+};