File size: 3,664 Bytes
246d201 |
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 |
import {
addAssistantMessage,
addAssistantAction,
addUserMessage,
addErrorMessage,
} from "#/state/chat-slice";
import { appendSecurityAnalyzerInput } from "#/state/security-analyzer-slice";
import { setCode, setActiveFilepath } from "#/state/code-slice";
import { appendJupyterInput } from "#/state/jupyter-slice";
import { setCurStatusMessage } from "#/state/status-slice";
import store from "#/store";
import ActionType from "#/types/action-type";
import {
ActionMessage,
ObservationMessage,
StatusMessage,
} from "#/types/message";
import { handleObservationMessage } from "./observations";
import { appendInput } from "#/state/command-slice";
const messageActions = {
[ActionType.BROWSE]: (message: ActionMessage) => {
if (!message.args.thought && message.message) {
store.dispatch(addAssistantMessage(message.message));
}
},
[ActionType.BROWSE_INTERACTIVE]: (message: ActionMessage) => {
if (!message.args.thought && message.message) {
store.dispatch(addAssistantMessage(message.message));
}
},
[ActionType.WRITE]: (message: ActionMessage) => {
const { path, content } = message.args;
store.dispatch(setActiveFilepath(path));
store.dispatch(setCode(content));
},
[ActionType.MESSAGE]: (message: ActionMessage) => {
if (message.source === "user") {
store.dispatch(
addUserMessage({
content: message.args.content,
imageUrls:
typeof message.args.image_urls === "string"
? [message.args.image_urls]
: message.args.image_urls,
timestamp: message.timestamp,
pending: false,
}),
);
} else {
store.dispatch(addAssistantMessage(message.args.content));
}
},
[ActionType.RUN_IPYTHON]: (message: ActionMessage) => {
if (message.args.confirmation_state !== "rejected") {
store.dispatch(appendJupyterInput(message.args.code));
}
},
};
export function handleActionMessage(message: ActionMessage) {
if (message.args?.hidden) {
return;
}
if (message.action === ActionType.RUN) {
store.dispatch(appendInput(message.args.command));
}
if ("args" in message && "security_risk" in message.args) {
store.dispatch(appendSecurityAnalyzerInput(message));
}
if (message.source === "agent") {
if (message.args && message.args.thought) {
store.dispatch(addAssistantMessage(message.args.thought));
}
// Need to convert ActionMessage to RejectAction
// @ts-expect-error TODO: fix
store.dispatch(addAssistantAction(message));
}
if (message.action in messageActions) {
const actionFn =
messageActions[message.action as keyof typeof messageActions];
actionFn(message);
}
}
export function handleStatusMessage(message: StatusMessage) {
if (message.type === "info") {
store.dispatch(
setCurStatusMessage({
...message,
}),
);
} else if (message.type === "error") {
store.dispatch(
addErrorMessage({
...message,
}),
);
}
}
export function handleAssistantMessage(message: Record<string, unknown>) {
if (message.action) {
handleActionMessage(message as unknown as ActionMessage);
} else if (message.observation) {
handleObservationMessage(message as unknown as ObservationMessage);
} else if (message.status_update) {
handleStatusMessage(message as unknown as StatusMessage);
} else {
store.dispatch(
addErrorMessage({
message: "Unknown message type received",
}),
);
}
}
|