Spaces:
Running
Running
File size: 3,267 Bytes
f5ed9bf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 |
import { EventEmitter, WebSocket } from 'ws';
import { BaseMessage, AIMessage, HumanMessage } from '@langchain/core/messages';
import handleWebSearch from '../agents/webSearchAgent';
import handleAcademicSearch from '../agents/academicSearchAgent';
import handleWritingAssistant from '../agents/writingAssistant';
import handleWolframAlphaSearch from '../agents/wolframAlphaSearchAgent';
import handleYoutubeSearch from '../agents/youtubeSearchAgent';
import handleRedditSearch from '../agents/redditSearchAgent';
import type { BaseChatModel } from '@langchain/core/language_models/chat_models';
import type { Embeddings } from '@langchain/core/embeddings';
import logger from '../utils/logger';
type Message = {
type: string;
content: string;
copilot: boolean;
focusMode: string;
history: Array<[string, string]>;
};
const searchHandlers = {
webSearch: handleWebSearch,
academicSearch: handleAcademicSearch,
writingAssistant: handleWritingAssistant,
wolframAlphaSearch: handleWolframAlphaSearch,
youtubeSearch: handleYoutubeSearch,
redditSearch: handleRedditSearch,
};
const handleEmitterEvents = (
emitter: EventEmitter,
ws: WebSocket,
id: string,
) => {
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 }));
});
};
export const handleMessage = async (
message: string,
ws: WebSocket,
llm: BaseChatModel,
embeddings: Embeddings,
) => {
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') {
const handler = searchHandlers[parsedMessage.focusMode];
if (handler) {
const emitter = handler(
parsedMessage.content,
history,
llm,
embeddings,
);
handleEmitterEvents(emitter, ws, id);
} else {
ws.send(JSON.stringify({ type: 'error', data: 'Invalid focus mode' }));
}
}
} catch (err) {
ws.send(JSON.stringify({ type: 'error', data: 'Invalid message format' }));
logger.error(`Failed to handle message: ${err}`);
}
};
|