Initial commit

src/websocket/connectionManager.ts

@@ -0,0 +1,11 @@
+import { WebSocket } from 'ws';
+import { handleMessage } from './messageHandler';
+
+export const handleConnection = (ws: WebSocket) => {
+  ws.on(
+    'message',
+    async (message) => await handleMessage(message.toString(), ws),
+  );
+
+  ws.on('close', () => console.log('Connection closed'));
+};

src/websocket/index.ts

@@ -0,0 +1,8 @@
+import { initServer } from './websocketServer';
+import http from 'http';
+
+export const startWebSocketServer = (
+  server: http.Server<typeof http.IncomingMessage, typeof http.ServerResponse>,
+) => {
+  initServer(server);
+};

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' }));
+  }
+};

src/websocket/websocketServer.ts

@@ -0,0 +1,15 @@
+import { WebSocketServer } from 'ws';
+import { handleConnection } from './connectionManager';
+import http from 'http';
+
+export const initServer = (
+  server: http.Server<typeof http.IncomingMessage, typeof http.ServerResponse>,
+) => {
+  const wss = new WebSocketServer({ server });
+
+  wss.on('connection', (ws) => {
+    handleConnection(ws);
+  });
+
+  console.log(`WebSocket server started on port ${process.env.PORT}`);
+};