diff --git a/frontend/src/components/variants/Variants.tsx b/frontend/src/components/variants/Variants.tsx
new file mode 100644
index 0000000..219841c
--- /dev/null
+++ b/frontend/src/components/variants/Variants.tsx
@@ -0,0 +1,64 @@
+import { useProjectStore } from "../../store/project-store";
+
+function Variants() {
+ const {
+ // Inputs
+ referenceImages,
+
+ // Outputs
+ variants,
+ currentVariantIndex,
+ setCurrentVariantIndex,
+ setGeneratedCode,
+ appHistory,
+ setAppHistory,
+ } = useProjectStore();
+
+ function switchVariant(index: number) {
+ const variant = variants[index];
+ setCurrentVariantIndex(index);
+ setGeneratedCode(variant);
+ if (appHistory.length === 1) {
+ setAppHistory([
+ {
+ type: "ai_create",
+ parentIndex: null,
+ code: variant,
+ inputs: { image_url: referenceImages[0] },
+ },
+ ]);
+ } else {
+ setAppHistory((prev) => {
+ const newHistory = [...prev];
+ newHistory[newHistory.length - 1].code = variant;
+ return newHistory;
+ });
+ }
+ }
+
+ if (variants.length === 0) {
+ return null;
+ }
+
+ return (
+
+
+ {variants.map((_, index) => (
+
switchVariant(index)}
+ >
+
Option {index + 1}
+
+ ))}
+
+
+ );
+}
+
+export default Variants;
diff --git a/frontend/src/generateCode.ts b/frontend/src/generateCode.ts
index 7fc34e7..b373702 100644
--- a/frontend/src/generateCode.ts
+++ b/frontend/src/generateCode.ts
@@ -11,12 +11,18 @@ const ERROR_MESSAGE =
const CANCEL_MESSAGE = "Code generation cancelled";
+type WebSocketResponse = {
+ type: "chunk" | "status" | "setCode" | "error";
+ value: string;
+ variantIndex: number;
+};
+
export function generateCode(
wsRef: React.MutableRefObject
,
params: FullGenerationSettings,
- onChange: (chunk: string) => void,
- onSetCode: (code: string) => void,
- onStatusUpdate: (status: string) => void,
+ onChange: (chunk: string, variantIndex: number) => void,
+ onSetCode: (code: string, variantIndex: number) => void,
+ onStatusUpdate: (status: string, variantIndex: number) => void,
onCancel: () => void,
onComplete: () => void
) {
@@ -31,13 +37,13 @@ export function generateCode(
});
ws.addEventListener("message", async (event: MessageEvent) => {
- const response = JSON.parse(event.data);
+ const response = JSON.parse(event.data) as WebSocketResponse;
if (response.type === "chunk") {
- onChange(response.value);
+ onChange(response.value, response.variantIndex);
} else if (response.type === "status") {
- onStatusUpdate(response.value);
+ onStatusUpdate(response.value, response.variantIndex);
} else if (response.type === "setCode") {
- onSetCode(response.value);
+ onSetCode(response.value, response.variantIndex);
} else if (response.type === "error") {
console.error("Error generating code", response.value);
toast.error(response.value);
diff --git a/frontend/src/store/project-store.ts b/frontend/src/store/project-store.ts
index 8a416f0..20516ae 100644
--- a/frontend/src/store/project-store.ts
+++ b/frontend/src/store/project-store.ts
@@ -16,10 +16,17 @@ interface ProjectStore {
setGeneratedCode: (
updater: string | ((currentCode: string) => string)
) => void;
- executionConsole: string[];
- setExecutionConsole: (
- updater: string[] | ((currentConsole: string[]) => string[])
- ) => void;
+
+ variants: string[];
+ currentVariantIndex: number;
+ setCurrentVariantIndex: (index: number) => void;
+ setVariant: (code: string, index: number) => void;
+ appendToVariant: (newTokens: string, index: number) => void;
+ resetVariants: () => void;
+
+ executionConsoles: { [key: number]: string[] };
+ appendExecutionConsole: (variantIndex: number, line: string) => void;
+ resetExecutionConsoles: () => void;
// Tracks the currently shown version from app history
// TODO: might want to move to appStore
@@ -48,14 +55,41 @@ export const useProjectStore = create((set) => ({
generatedCode:
typeof updater === "function" ? updater(state.generatedCode) : updater,
})),
- executionConsole: [],
- setExecutionConsole: (updater) =>
+
+ variants: [],
+ currentVariantIndex: 0,
+
+ setCurrentVariantIndex: (index) => set({ currentVariantIndex: index }),
+ setVariant: (code: string, index: number) =>
+ set((state) => {
+ const newVariants = [...state.variants];
+ while (newVariants.length <= index) {
+ newVariants.push("");
+ }
+ newVariants[index] = code;
+ return { variants: newVariants };
+ }),
+ appendToVariant: (newTokens: string, index: number) =>
+ set((state) => {
+ const newVariants = [...state.variants];
+ newVariants[index] += newTokens;
+ return { variants: newVariants };
+ }),
+ resetVariants: () => set({ variants: [], currentVariantIndex: 0 }),
+
+ executionConsoles: {},
+
+ appendExecutionConsole: (variantIndex: number, line: string) =>
set((state) => ({
- executionConsole:
- typeof updater === "function"
- ? updater(state.executionConsole)
- : updater,
+ executionConsoles: {
+ ...state.executionConsoles,
+ [variantIndex]: [
+ ...(state.executionConsoles[variantIndex] || []),
+ line,
+ ],
+ },
})),
+ resetExecutionConsoles: () => set({ executionConsoles: {} }),
currentVersion: null,
setCurrentVersion: (version) => set({ currentVersion: version }),