Personal-Test-Ground/src/FFB/ffb.gateway.ts

import {
    WebSocketGateway,
    WebSocketServer,
    OnGatewayInit,
    OnGatewayConnection,
    OnGatewayDisconnect,
    SubscribeMessage,
    MessageBody,
    ConnectedSocket,
} from '@nestjs/websockets';
import { Server, Socket } from 'socket.io';
import { Logger, Inject, forwardRef } from '@nestjs/common';
import { FFBLangChainService } from './services/ffb-langchain.service';

@WebSocketGateway({
    cors: {
        origin: '*',
    },
    namespace: 'ffb',
})
export class FFBGateway
    implements OnGatewayInit, OnGatewayConnection, OnGatewayDisconnect {
    @WebSocketServer() server: Server;
    private logger: Logger = new Logger('FFBGateway');

    constructor(
        @Inject(forwardRef(() => FFBLangChainService))
        private readonly langchainService: FFBLangChainService
    ) { }

    afterInit(server: Server) {
        this.logger.log('FFB Gateway Initialized');
    }

    handleConnection(client: Socket, ...args: any[]) {
        this.logger.log(`Client connected: ${client.id}`);
        this.langchainService.createSession(client.id);
    }

    handleDisconnect(client: Socket) {
        this.logger.log(`Client disconnected: ${client.id}`);
        this.langchainService.deleteSession(client.id);
    }

    @SubscribeMessage('chat')
    async handleChat(
        @MessageBody() data: { message: string },
        @ConnectedSocket() client: Socket,
    ) {
        this.logger.log(`Received chat from ${client.id}: ${data.message}`);
        const response = await this.langchainService.chat(client.id, data.message);
        client.emit('chat_response', { message: response });
        return { event: 'sent', data: response }; // Acknowledgement if needed
    }

    @SubscribeMessage('switch_model')
    handleSwitchModel(
        @MessageBody() data: { provider: 'openai' | 'gemini' },
        @ConnectedSocket() client: Socket,
    ) {
        this.logger.log(`Switching model for ${client.id} to ${data.provider}`);
        this.langchainService.switchModel(client.id, data.provider);
    }

    @SubscribeMessage('get_model')
    handleGetModel(@ConnectedSocket() client: Socket) {
        const model = this.langchainService.getCurrentModel(client.id);
        console.log(model)
        return { event: 'current_model', data: model };
    }

    emitThought(socketId: string, data: any) {
        this.server.to(socketId).emit('agent_thought', data);
    }
}