openreplay/frontend/app/components/Kai/SocketManager.ts
2025-05-13 13:58:11 +02:00

92 lines
2.3 KiB
TypeScript

import io from 'socket.io-client';
export class ChatManager {
socket: ReturnType<typeof io>;
threadId: string | null = null;
constructor({ projectId, threadId, token }: { projectId: string; threadId: string, token: string }) {
this.threadId = threadId;
console.log('Kai socket', projectId, threadId, token);
const socket = io(`localhost:8700/kai/chat`, {
transports: ['websocket'],
autoConnect: true,
reconnection: true,
reconnectionAttempts: 5,
reconnectionDelay: 1000,
reconnectionDelayMax: 5000,
withCredentials: true,
multiplex: true,
query: {
project_id: projectId,
thread_id: threadId,
},
auth: {
token: `Bearer ${token}`,
}
});
socket.on('connect', () => {
console.log('Connected to server');
});
socket.on('disconnect', () => {
console.log('Disconnected from server');
});
socket.on('error', (err) => {
console.error('Socket error:', err);
})
socket.onAny((e) => console.log('event', e));
this.socket = socket;
}
sendMessage = (message: string, isReplace = false) => {
this.socket.emit(
'message',
JSON.stringify({
message,
threadId: this.threadId,
replace: isReplace,
}),
);
};
setOnMsgHook = ({
msgCallback,
titleCallback,
}: {
msgCallback: (msg: BotChunk | { state: string, type: 'state', start_time?: number }) => void;
titleCallback: (title: string) => void;
}) => {
this.socket.on('chunk', (msg: BotChunk) => {
console.log('Received message:', msg);
msgCallback(msg);
});
this.socket.on('title', (msg: { content: string }) => {
console.log('Received title:', msg);
titleCallback(msg.content);
});
this.socket.on('state', (state: { message: 'idle' | 'running', start_time: number }) => {
msgCallback({ state: state.message, type: 'state', start_time: state.start_time })
})
};
disconnect = () => {
this.socket.disconnect();
};
}
export interface BotChunk {
stage: 'start' | 'chart' | 'final' | 'title';
content: string;
messageId: string;
duration?: number;
}
export interface Message {
text: string;
isUser: boolean;
messageId: string;
duration?: number;
}
export interface SentMessage extends Message {
replace: boolean
}