diff --git a/src/IOC/ioc.config.ts b/src/IOC/ioc.config.ts index aa3dce7..0eb7b51 100644 --- a/src/IOC/ioc.config.ts +++ b/src/IOC/ioc.config.ts @@ -58,6 +58,7 @@ import { ChatbotService } from "../modules/chatbot/providers/chatbot.service"; import { DataContextService } from "../modules/chatbot/providers/data-context.service"; import { LLMService } from "../modules/chatbot/providers/llm.service"; import { WebSocketAuthService } from "../modules/chatbot/providers/websocket-auth.service"; +import { ChatbotGateway } from "../modules/chatbot/chatbot.gateway"; import { ChatSessionRepository, createChatSessionRepository } from "../modules/chatbot/repositories/chat-session.repository"; import { ChatMessageRepository, createChatMessageRepository } from "../modules/chatbot/repositories/chat-message.repository"; import { ContactUsRepo, CreateContactUsRepo } from "../modules/contact-us/contactUs.repository"; @@ -207,6 +208,7 @@ const containerModules = new AsyncContainerModule(async (bind) => { bind(IOCTYPES.ChatbotLLMService).to(LLMService).inSingletonScope(); bind(IOCTYPES.ChatbotDataContextService).to(DataContextService).inSingletonScope(); bind(IOCTYPES.ChatbotWebSocketAuthService).to(WebSocketAuthService).inSingletonScope(); + bind(IOCTYPES.ChatbotGateway).to(ChatbotGateway).inSingletonScope(); // #endregion // #region repository bind(IOCTYPES.CategoryRepository).toDynamicValue(createCategoryRepo).inSingletonScope(); diff --git a/src/modules/chatbot/chatbot.gateway.ts b/src/modules/chatbot/chatbot.gateway.ts new file mode 100644 index 0000000..09ab37f --- /dev/null +++ b/src/modules/chatbot/chatbot.gateway.ts @@ -0,0 +1,330 @@ +import { Server as HttpServer } from "http"; + +import { inject, injectable } from "inversify"; +import { Server, Socket } from "socket.io"; + +import { ChatbotService } from "./providers/chatbot.service"; +import { WebSocketAuthService } from "./providers/websocket-auth.service"; +import { SendMessageDto } from "./DTO/send-message.dto"; +import { WEBSOCKET_EVENTS } from "./constants/chatbot.constants"; +import { IOCTYPES } from "../../IOC/ioc.types"; +import { Logger } from "../../core/logging/logger"; +import { AuthenticatedSocket, WebSocketResponse } from "./interfaces/websocket.interface"; +import * as ulidLib from "ulid"; + +@injectable() +export class ChatbotGateway { + private io: Server; + private readonly logger: Logger; + private readonly CHATBOT_ULID_COOKIE_NAME = "chatbot_session_id"; + + constructor( + @inject(IOCTYPES.ChatbotService) private chatbotService: ChatbotService, + @inject(IOCTYPES.ChatbotWebSocketAuthService) private wsAuthService: WebSocketAuthService, + ) { + this.logger = new Logger("ChatbotGateway"); + } + + public initialize(server: HttpServer): Server { + this.io = new Server(server, { + path: "/ws-chatbot", + cors: { + origin: true, + allowedHeaders: ["Authorization"], + credentials: true, + methods: ["GET", "POST"], + }, + }); + + this.io.on("connection", (socket: Socket) => this.handleConnection(socket)); + this.logger.info("ChatbotGateway initialized on /ws-chatbot"); + return this.io; + } + + private async handleConnection(socket: Socket): Promise { + try { + // Try to authenticate if token is provided + const token = socket.handshake.query.token as string; + let ulid: string | undefined; + + if (token) { + // Authenticated connection + const authResult = await this.wsAuthService.authenticateClient(socket as AuthenticatedSocket); + if (!authResult.success) { + this.wsAuthService.handleAuthenticationFailure(socket, authResult.error || "Authentication failed"); + return; + } + this.wsAuthService.emitAuthenticationSuccess(socket, authResult.user!); + ulid = authResult.user!.id; + } else { + // Anonymous connection - generate or get ULID from cookie + const cookieHeader = socket.handshake.headers.cookie; + ulid = this.extractUlidFromCookie(cookieHeader); + if (!ulid) { + ulid = ulidLib.ulid(); + } + socket.data = { user: { id: ulid, sub: ulid } }; + } + + this.logger.info(`Client connected: ${socket.id}, ULID: ${ulid}`); + + // Register event handlers + this.registerEventHandlers(socket, ulid); + } catch (error) { + this.logger.error("Connection error", error); + socket.emit(WEBSOCKET_EVENTS.ERROR, { + status: "error", + message: "Connection failed", + timestamp: new Date().toISOString(), + } as WebSocketResponse); + socket.disconnect(); + } + } + + private extractUlidFromCookie(cookieHeader?: string | string[]): string | undefined { + if (!cookieHeader || typeof cookieHeader !== "string") { + return undefined; + } + + const cookies = cookieHeader.split(";").reduce( + (acc, cookie) => { + const [key, value] = cookie.trim().split("="); + if (key && value) { + acc[key] = decodeURIComponent(value); + } + return acc; + }, + {} as Record, + ); + + return cookies[this.CHATBOT_ULID_COOKIE_NAME]; + } + + private registerEventHandlers(socket: Socket, ulid: string): void { + // Session management + socket.on(WEBSOCKET_EVENTS.CREATE_SESSION, () => this.handleCreateSession(socket, ulid)); + socket.on(WEBSOCKET_EVENTS.JOIN_CHAT, (sessionId: string) => this.handleJoinChat(socket, ulid, sessionId)); + socket.on(WEBSOCKET_EVENTS.LEAVE_CHAT, (sessionId: string) => this.handleLeaveChat(socket, sessionId)); + + // Message handling + socket.on(WEBSOCKET_EVENTS.SEND_MESSAGE, (data: SendMessageDto) => this.handleSendMessage(socket, ulid, data)); + socket.on(WEBSOCKET_EVENTS.SEND_MESSAGE + "_stream", (data: SendMessageDto) => this.handleSendMessageStream(socket, ulid, data)); + + // Typing indicators + socket.on(WEBSOCKET_EVENTS.TYPING_START, (sessionId: string) => this.handleTypingStart(socket, sessionId)); + socket.on(WEBSOCKET_EVENTS.TYPING_STOP, (sessionId: string) => this.handleTypingStop(socket, sessionId)); + + // Disconnect + socket.on(WEBSOCKET_EVENTS.DISCONNECT, () => this.handleDisconnect(socket)); + } + + private async handleCreateSession(socket: Socket, ulid: string): Promise { + try { + const session = await this.chatbotService.createChatSession(ulid); + socket.emit(WEBSOCKET_EVENTS.SESSION_CREATED, { + status: "success", + data: session, + timestamp: new Date().toISOString(), + } as WebSocketResponse); + } catch (error) { + this.logger.error("Error creating session", error); + socket.emit(WEBSOCKET_EVENTS.SESSION_ERROR, { + status: "error", + message: "Failed to create session", + timestamp: new Date().toISOString(), + } as WebSocketResponse); + } + } + + private async handleJoinChat(socket: Socket, ulid: string, sessionId: string): Promise { + try { + const session = await this.chatbotService.getChatSession(sessionId, ulid); + socket.join(`session_${sessionId}`); + (socket as AuthenticatedSocket).data.sessionId = sessionId; + + socket.emit(WEBSOCKET_EVENTS.CHAT_JOINED, { + status: "success", + data: session, + timestamp: new Date().toISOString(), + } as WebSocketResponse); + + // Notify others in the session (if any) + socket.to(`session_${sessionId}`).emit(WEBSOCKET_EVENTS.USER_JOINED, { + status: "success", + data: { userId: ulid, sessionId }, + timestamp: new Date().toISOString(), + } as WebSocketResponse); + } catch (error) { + this.logger.error("Error joining chat", error); + socket.emit(WEBSOCKET_EVENTS.SESSION_ERROR, { + status: "error", + message: "Failed to join chat session", + timestamp: new Date().toISOString(), + } as WebSocketResponse); + } + } + + private handleLeaveChat(socket: Socket, sessionId: string): void { + socket.leave(`session_${sessionId}`); + delete (socket as AuthenticatedSocket).data.sessionId; + + socket.emit(WEBSOCKET_EVENTS.CHAT_LEFT, { + status: "success", + message: "Left chat session", + timestamp: new Date().toISOString(), + } as WebSocketResponse); + + socket.to(`session_${sessionId}`).emit(WEBSOCKET_EVENTS.USER_LEFT, { + status: "success", + data: { sessionId }, + timestamp: new Date().toISOString(), + } as WebSocketResponse); + } + + private async handleSendMessage(socket: Socket, ulid: string, data: SendMessageDto): Promise { + try { + // Save user message + const userMessage = await this.chatbotService.sendMessage(ulid, data); + + // Emit user message to client + socket.emit(WEBSOCKET_EVENTS.MESSAGE_RECEIVED, { + status: "success", + data: userMessage, + timestamp: new Date().toISOString(), + } as WebSocketResponse); + + // Emit to others in the session (if any) + socket.to(`session_${data.sessionId}`).emit(WEBSOCKET_EVENTS.MESSAGE_RECEIVED, { + status: "success", + data: userMessage, + timestamp: new Date().toISOString(), + } as WebSocketResponse); + + // Generate bot response asynchronously + this.generateBotResponseAsync(socket, data.sessionId, userMessage.id, ulid); + } catch (error) { + this.logger.error("Error sending message", error); + socket.emit(WEBSOCKET_EVENTS.MESSAGE_ERROR, { + status: "error", + message: "Failed to send message", + timestamp: new Date().toISOString(), + } as WebSocketResponse); + } + } + + private async handleSendMessageStream(socket: Socket, ulid: string, data: SendMessageDto): Promise { + try { + const { userMessage, streamGenerator } = await this.chatbotService.sendMessageStream(ulid, data); + + // Emit user message + socket.emit(WEBSOCKET_EVENTS.MESSAGE_RECEIVED, { + status: "success", + data: userMessage, + timestamp: new Date().toISOString(), + } as WebSocketResponse); + + // Emit to others in the session (if any) + socket.to(`session_${data.sessionId}`).emit(WEBSOCKET_EVENTS.MESSAGE_RECEIVED, { + status: "success", + data: userMessage, + timestamp: new Date().toISOString(), + } as WebSocketResponse); + + // Start streaming bot response + socket.emit(WEBSOCKET_EVENTS.BOT_RESPONSE_START, { + status: "success", + data: { userMessageId: userMessage.id }, + timestamp: new Date().toISOString(), + } as WebSocketResponse); + + let fullResponse = ""; + const stream = await streamGenerator(); + for await (const chunk of stream) { + if (chunk) { + fullResponse += chunk; + socket.emit(WEBSOCKET_EVENTS.BOT_RESPONSE_CHUNK, { + status: "success", + data: { chunk, userMessageId: userMessage.id }, + timestamp: new Date().toISOString(), + } as WebSocketResponse); + } + } + + // Save bot response after streaming completes + if (fullResponse) { + await this.chatbotService.saveStreamedBotResponse(data.sessionId, userMessage.id, fullResponse, ulid); + } + + // Emit final bot response + socket.emit(WEBSOCKET_EVENTS.BOT_RESPONSE_END, { + status: "success", + data: { userMessageId: userMessage.id, fullResponse }, + timestamp: new Date().toISOString(), + } as WebSocketResponse); + } catch (error) { + this.logger.error("Error sending stream message", error); + socket.emit(WEBSOCKET_EVENTS.MESSAGE_ERROR, { + status: "error", + message: "Failed to send message", + timestamp: new Date().toISOString(), + } as WebSocketResponse); + } + } + + private async generateBotResponseAsync(socket: Socket, sessionId: string, userMessageId: string, ulid: string): Promise { + try { + // Get the session to check if it exists + const session = await this.chatbotService.getChatSession(sessionId, ulid); + if (!session) { + return; + } + + // The bot response will be generated asynchronously by the service + // We'll need to emit it when it's ready + // For now, we'll rely on the service's internal async generation + // In a production system, you might want to use Redis pub/sub or similar + // to notify when the bot response is ready + + // Emit that bot is typing + socket.emit(WEBSOCKET_EVENTS.BOT_RESPONSE_START, { + status: "success", + data: { userMessageId }, + timestamp: new Date().toISOString(), + } as WebSocketResponse); + + // Note: The actual bot response will be generated by ChatbotService.generateBotResponse + // which runs asynchronously. In a real implementation, you might want to use + // Redis pub/sub or a message queue to notify when the response is ready + } catch (error) { + this.logger.error("Error generating bot response", error); + socket.emit(WEBSOCKET_EVENTS.BOT_RESPONSE, { + status: "error", + message: "Failed to generate bot response", + timestamp: new Date().toISOString(), + } as WebSocketResponse); + } + } + + private handleTypingStart(socket: Socket, sessionId: string): void { + socket.to(`session_${sessionId}`).emit(WEBSOCKET_EVENTS.TYPING_START, { + status: "success", + data: { sessionId, userId: socket.data?.user?.id }, + timestamp: new Date().toISOString(), + } as WebSocketResponse); + } + + private handleTypingStop(socket: Socket, sessionId: string): void { + socket.to(`session_${sessionId}`).emit(WEBSOCKET_EVENTS.TYPING_STOP, { + status: "success", + data: { sessionId, userId: socket.data?.user?.id }, + timestamp: new Date().toISOString(), + } as WebSocketResponse); + } + + private handleDisconnect(socket: Socket): void { + const user = socket.data?.user; + if (user) { + this.logger.info(`Client disconnected: ${socket.id}, ULID: ${user.id}`); + } + } +} diff --git a/src/server.ts b/src/server.ts index e0b438d..490ff36 100644 --- a/src/server.ts +++ b/src/server.ts @@ -11,6 +11,7 @@ import { Logger } from "./core/logging/logger"; import { connectMongo } from "./db/connection"; import { IOCTYPES } from "./IOC/ioc.types"; import { ChatGateway } from "./modules/chat/chat.gateway"; +import { ChatbotGateway } from "./modules/chatbot/chatbot.gateway"; import { StartWorker } from "./queues"; import { gracefulShutdown } from "./utils/shutdown.utils"; @@ -28,6 +29,9 @@ async function bootStrap() { const chatGateway = container.get(IOCTYPES.ChatGateway); const wsServer = chatGateway.initialize(server); + const chatbotGateway = container.get(IOCTYPES.ChatbotGateway); + const chatbotWsServer = chatbotGateway.initialize(server); + const worker = await StartWorker(); // @@ -36,9 +40,9 @@ async function bootStrap() { logger.info(`swagger documentation is serving on http://localhost:${PORT}/api-docs`); }); - process.on("SIGINT", () => gracefulShutdown("SIGINT", server, wsServer, worker[0], logger)); + process.on("SIGINT", () => gracefulShutdown("SIGINT", server, wsServer, worker[0], logger, chatbotWsServer)); - process.on("SIGTERM", () => gracefulShutdown("SIGTERM", server, wsServer, worker[0], logger)); + process.on("SIGTERM", () => gracefulShutdown("SIGTERM", server, wsServer, worker[0], logger, chatbotWsServer)); } catch (error) { logger.error("Error starting the server.", error); process.exit(1); diff --git a/src/utils/shutdown.utils.ts b/src/utils/shutdown.utils.ts index 51d4ff0..a100e51 100644 --- a/src/utils/shutdown.utils.ts +++ b/src/utils/shutdown.utils.ts @@ -5,16 +5,29 @@ import { Server as WSServer } from "socket.io"; import { Logger } from "../core/logging/logger"; -export const gracefulShutdown = async (signal: string, server: Server, wsServer: WSServer, worker: Worker, logger: Logger) => { +export const gracefulShutdown = async ( + signal: string, + server: Server, + wsServer: WSServer, + worker: Worker, + logger: Logger, + chatbotWsServer?: WSServer, +) => { try { logger.warn(`${signal} signal received: closing HTTP server and worker...`); logger.warn("close all socket connection.."); wsServer.disconnectSockets(); + if (chatbotWsServer) { + chatbotWsServer.disconnectSockets(); + } logger.warn("all socket connection disconnected"); logger.warn("Closing ws server..."); wsServer.close(); + if (chatbotWsServer) { + chatbotWsServer.close(); + } logger.warn("ws server closed"); server.close(async () => {