File size: 4,885 Bytes
f8b5d42 |
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 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 |
const chalk = require("chalk");
const { Telemetry } = require("../../../../models/telemetry");
const SOCKET_TIMEOUT_MS = 300 * 1_000; // 5 mins
/**
* Websocket Interface plugin. It prints the messages on the console and asks for feedback
* while the conversation is running in the background.
*/
// export interface AIbitatWebSocket extends ServerWebSocket<unknown> {
// askForFeedback?: any
// awaitResponse?: any
// handleFeedback?: (message: string) => void;
// }
const WEBSOCKET_BAIL_COMMANDS = [
"exit",
"/exit",
"stop",
"/stop",
"halt",
"/halt",
"/reset", // Will not reset but will bail. Powerusers always do this and the LLM responds.
];
const websocket = {
name: "websocket",
startupConfig: {
params: {
socket: {
required: true,
},
muteUserReply: {
required: false,
default: true,
},
introspection: {
required: false,
default: true,
},
},
},
plugin: function ({
socket, // @type AIbitatWebSocket
muteUserReply = true, // Do not post messages to "USER" back to frontend.
introspection = false, // when enabled will attach socket to Aibitat object with .introspect method which reports status updates to frontend.
}) {
return {
name: this.name,
setup(aibitat) {
aibitat.onError(async (error) => {
let errorMessage =
error?.message || "An error occurred while running the agent.";
console.error(chalk.red(` error: ${errorMessage}`), error);
aibitat.introspect(
`Error encountered while running: ${errorMessage}`
);
socket.send(
JSON.stringify({ type: "wssFailure", content: errorMessage })
);
aibitat.terminate();
});
aibitat.introspect = (messageText) => {
if (!introspection) return; // Dump thoughts when not wanted.
socket.send(
JSON.stringify({
type: "statusResponse",
content: messageText,
animate: true,
})
);
};
// expose function for sockets across aibitat
// type param must be set or else msg will not be shown or handled in UI.
aibitat.socket = {
send: (type = "__unhandled", content = "") => {
socket.send(JSON.stringify({ type, content }));
},
};
// aibitat.onStart(() => {
// console.log("🚀 starting chat ...");
// });
aibitat.onMessage((message) => {
if (message.from !== "USER")
Telemetry.sendTelemetry("agent_chat_sent");
if (message.from === "USER" && muteUserReply) return;
socket.send(JSON.stringify(message));
});
aibitat.onTerminate(() => {
// console.log("🚀 chat finished");
socket.close();
});
aibitat.onInterrupt(async (node) => {
const feedback = await socket.askForFeedback(socket, node);
if (WEBSOCKET_BAIL_COMMANDS.includes(feedback)) {
socket.close();
return;
}
await aibitat.continue(feedback);
});
/**
* Socket wait for feedback on socket
*
* @param socket The content to summarize. // AIbitatWebSocket & { receive: any, echo: any }
* @param node The chat node // { from: string; to: string }
* @returns The summarized content.
*/
socket.askForFeedback = (socket, node) => {
socket.awaitResponse = (question = "waiting...") => {
socket.send(JSON.stringify({ type: "WAITING_ON_INPUT", question }));
return new Promise(function (resolve) {
let socketTimeout = null;
socket.handleFeedback = (message) => {
const data = JSON.parse(message);
if (data.type !== "awaitingFeedback") return;
delete socket.handleFeedback;
clearTimeout(socketTimeout);
resolve(data.feedback);
return;
};
socketTimeout = setTimeout(() => {
console.log(
chalk.red(
`Client took too long to respond, chat thread is dead after ${SOCKET_TIMEOUT_MS}ms`
)
);
resolve("exit");
return;
}, SOCKET_TIMEOUT_MS);
});
};
return socket.awaitResponse(`Provide feedback to ${chalk.yellow(
node.to
)} as ${chalk.yellow(node.from)}.
Press enter to skip and use auto-reply, or type 'exit' to end the conversation: \n`);
};
// console.log("🚀 WS plugin is complete.");
},
};
},
};
module.exports = {
websocket,
WEBSOCKET_BAIL_COMMANDS,
};
|