anything-llm/server/endpoints/agentWebsocket.js

62 lines
1.8 KiB
JavaScript
Raw Normal View History

const { Telemetry } = require("../models/telemetry");
const {
WorkspaceAgentInvocation,
} = require("../models/workspaceAgentInvocation");
const { AgentHandler } = require("../utils/agents");
const {
WEBSOCKET_BAIL_COMMANDS,
} = require("../utils/agents/aibitat/plugins/websocket");
const { safeJsonParse } = require("../utils/http");
// Setup listener for incoming messages to relay to socket so it can be handled by agent plugin.
function relayToSocket(message) {
if (this.handleFeedback) return this?.handleFeedback?.(message);
this.checkBailCommand(message);
}
function agentWebsocket(app) {
if (!app) return;
app.ws("/agent-invocation/:uuid", async function (socket, request) {
try {
const agentHandler = await new AgentHandler({
uuid: String(request.params.uuid),
}).init();
if (!agentHandler.invocation) {
socket.close();
return;
}
socket.on("message", relayToSocket);
socket.on("close", () => {
agentHandler.closeAlert();
WorkspaceAgentInvocation.close(String(request.params.uuid));
return;
});
socket.checkBailCommand = (data) => {
const content = safeJsonParse(data)?.feedback;
if (WEBSOCKET_BAIL_COMMANDS.includes(content)) {
agentHandler.log(
`User invoked bail command while processing. Closing session now.`
);
agentHandler.aibitat.abort();
socket.close();
return;
}
};
await Telemetry.sendTelemetry("agent_chat_started");
await agentHandler.createAIbitat({ socket });
await agentHandler.startAgentCluster();
} catch (e) {
console.error(e.message);
socket?.send(JSON.stringify({ type: "wssFailure", content: e.message }));
socket?.close();
}
});
}
module.exports = { agentWebsocket };