Code node generator built in to Rivet

packages/app/src/components/RivetApp.tsx

@@ -75,7 +75,15 @@ export const RivetApp: FC = () => {
       <DataStudioRenderer />
       <PluginsOverlayRenderer />
       <UpdateModalRenderer />
-      <ToastContainer position="bottom-right" hideProgressBar newestOnTop />
+      <ToastContainer enableMultiContainer position="bottom-right" hideProgressBar newestOnTop />
+      <ToastContainer
+        enableMultiContainer
+        containerId="wide"
+        style={{ width: 600 }}
+        position="bottom-right"
+        hideProgressBar
+        newestOnTop
+      />
     </div>
   );
 };

packages/app/src/components/editors/CodeEditor.tsx

@@ -34,14 +34,17 @@ export const DefaultCodeEditor: FC<
       const currentValue = (nodeLatest.current?.data as Record<string, unknown> | undefined)?.[editorDef.dataKey] as
         | string
         | undefined;
-      editorInstance.current.setValue(currentValue ?? '');
+
+      if (editorInstance.current.getValue() !== currentValue) {
+        editorInstance.current.setValue(currentValue ?? '');
+      }
 
       editorInstance.current.updateOptions({
         readOnly: isReadonly,
       });
     }
     // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [node.id, isReadonly]);
+  }, [node.id, isReadonly, (node.data as Record<string, unknown>)[editorDef.dataKey]]);
 
   const handleKeyDown = (e: monaco.IKeyboardEvent) => {
     if (e.keyCode === 9 /* Escape */) {

packages/app/src/components/editors/CustomEditor.tsx

@@ -0,0 +1,15 @@
+import { type CustomEditorDefinition, type ChartNode } from '@ironclad/rivet-core';
+import { type FC } from 'react';
+import { type SharedEditorProps } from './SharedEditorProps';
+import { match } from 'ts-pattern';
+import { CodeNodeAIAssistEditor } from './custom/CodeNodeAIAssistEditor';
+
+export const CustomEditor: FC<
+  SharedEditorProps & {
+    editor: CustomEditorDefinition<ChartNode>;
+  }
+> = ({ editor, ...props }) => {
+  return match(editor.customEditorId)
+    .with('CodeNodeAIAssist', () => <CodeNodeAIAssistEditor {...props} editor={editor} />)
+    .otherwise(() => null);
+};

packages/app/src/components/editors/DefaultNodeEditorField.tsx

@@ -20,6 +20,7 @@ import { DefaultToggleEditor } from './ToggleEditor';
 import { EditorGroup } from './EditorGroup';
 import { KeyValuePairEditor } from './KeyValuePairEditor';
 import { StringListEditor } from './StringListEditor';
+import { CustomEditor } from './CustomEditor';
 
 export const DefaultNodeEditorField: FC<
   SharedEditorProps & {
@@ -56,6 +57,7 @@ export const DefaultNodeEditorField: FC<
     .with({ type: 'group' }, (editor) => <EditorGroup {...sharedProps} editor={editor} />)
     .with({ type: 'keyValuePair' }, (editor) => <KeyValuePairEditor {...sharedProps} editor={editor} />)
     .with({ type: 'stringList' }, (editor) => <StringListEditor {...sharedProps} editor={editor} />)
+    .with({ type: 'custom' }, (editor) => <CustomEditor {...sharedProps} editor={editor} />)
     .exhaustive();
 
   const toggle =

packages/app/src/components/editors/custom/CodeNodeAIAssistEditor.tsx

@@ -0,0 +1,129 @@
+import { useState, type FC, useMemo } from 'react';
+import { type SharedEditorProps } from '../SharedEditorProps';
+import {
+  getError,
+  type ChartNode,
+  type CustomEditorDefinition,
+  coreCreateProcessor,
+  deserializeProject,
+  type CodeNodeData,
+  coerceType,
+  coerceTypeOptional,
+} from '@ironclad/rivet-core';
+import { Field } from '@atlaskit/form';
+import TextField from '@atlaskit/textfield';
+import Button from '@atlaskit/button';
+import { css } from '@emotion/react';
+import Select from '@atlaskit/select';
+import { toast } from 'react-toastify';
+import codeGeneratorProject from '../../../../graphs/code-node-generator.rivet-project?raw';
+import { useRecoilValue } from 'recoil';
+import { settingsState } from '../../../state/settings';
+import { fillMissingSettingsFromEnvironmentVariables } from '../../../utils/tauri';
+import { useDependsOnPlugins } from '../../../hooks/useDependsOnPlugins';
+import { marked } from 'marked';
+
+const styles = css`
+  display: flex;
+  align-items: center;
+  gap: 8px;
+
+  .model-selector {
+    width: 200px;
+  }
+`;
+
+export const CodeNodeAIAssistEditor: FC<
+  SharedEditorProps & {
+    editor: CustomEditorDefinition<ChartNode>;
+  }
+> = ({ node, isReadonly, isDisabled, onChange, editor }) => {
+  const [prompt, setPrompt] = useState('');
+  const [working, setWorking] = useState(false);
+  const [model, setModel] = useState('gpt-3.5-turbo');
+
+  const settings = useRecoilValue(settingsState);
+  const plugins = useDependsOnPlugins();
+
+  const data = node.data as CodeNodeData;
+
+  const generateCode = async () => {
+    try {
+      const [project] = deserializeProject(codeGeneratorProject);
+      const processor = coreCreateProcessor(project, {
+        inputs: {
+          prompt,
+          model,
+        },
+        ...(await fillMissingSettingsFromEnvironmentVariables(settings, plugins)),
+      });
+
+      setWorking(true);
+
+      const outputs = await processor.run();
+      const code = coerceTypeOptional(outputs.code, 'string');
+
+      if (code) {
+        onChange({
+          ...node,
+          data: {
+            ...data,
+            code,
+          } satisfies CodeNodeData,
+        });
+      } else {
+        const markdownResponse = marked(coerceType(outputs.response, 'string'));
+        toast.info(<div dangerouslySetInnerHTML={{ __html: markdownResponse }}></div>, {
+          autoClose: false,
+          containerId: 'wide',
+          toastId: 'ai-assist-response',
+        });
+      }
+    } catch (err) {
+      toast.error(`Failed to generate code: ${getError(err).message}`);
+    } finally {
+      setWorking(false);
+    }
+  };
+
+  const modelOptions = useMemo(
+    () => [
+      { label: 'GPT-4', value: 'gpt-4' },
+      { label: 'GPT-3.5 Turbo', value: 'gpt-3.5-turbo' },
+    ],
+    [],
+  );
+
+  const selectedModel = modelOptions.find((option) => option.value === model);
+
+  return (
+    <Field name="aiAssist" label="Generate Using AI">
+      {() => (
+        <div css={styles}>
+          <TextField
+            isDisabled={isDisabled || working}
+            isReadOnly={isReadonly}
+            value={prompt}
+            onChange={(e) => setPrompt((e.target as HTMLInputElement).value)}
+            placeholder="What would you like your code node to do?"
+            onKeyDown={(e) => {
+              if (e.key === 'Enter') {
+                generateCode();
+              }
+            }}
+          />
+          <Select
+            options={modelOptions}
+            value={selectedModel}
+            onChange={(option) => setModel(option!.value)}
+            isDisabled={isDisabled || working}
+            className="model-selector"
+          />
+          <Button appearance="primary" onClick={generateCode} isDisabled={isDisabled || working}>
+            Generate
+          </Button>
+        </div>
+      )}
+    </Field>
+  );
+};

packages/app/src/hooks/useSyncCurrentStateIntoOpenedProjects.ts

@@ -6,15 +6,15 @@ import {
   openedProjectsState,
   projectState,
 } from '../state/savedGraphs';
-import { graphMetadataState } from '../state/graph';
+import { graphState } from '../state/graph';
 
 export function useSyncCurrentStateIntoOpenedProjects() {
   const [openedProjects, setOpenedProjects] = useRecoilState(openedProjectsState);
   const [openedProjectsSortedIds, setOpenedProjectsSortedIds] = useRecoilState(openedProjectsSortedIdsState);
 
   const currentProject = useRecoilValue(projectState);
   const loadedProject = useRecoilValue(loadedProjectState);
-  const currentGraph = useRecoilValue(graphMetadataState);
+  const currentGraph = useRecoilValue(graphState);
 
   // Make sure current opened project is in opened projects
   useEffect(() => {
@@ -56,15 +56,24 @@ export function useSyncCurrentStateIntoOpenedProjects() {
 
   // Sync current graph into opened projects
   useEffect(() => {
-    if (currentGraph?.id != null && currentProject.graphs[currentGraph.id]) {
+    if (currentGraph.metadata?.id != null && currentProject.graphs[currentGraph.metadata.id]) {
       setOpenedProjects((openedProjects) => ({
         ...openedProjects,
         [currentProject.metadata.id]: {
           ...openedProjects[currentProject.metadata.id],
-          openedGraph: currentGraph!.id!,
+          project: {
+            ...currentProject,
+            graphs: {
+              // Sync current graph into opened projects as well, so that when you make changes, nav away, nav back, your changes are still there
+              ...currentProject.graphs,
+              [currentGraph.metadata!.id!]: currentGraph,
+            },
+          },
+          openedGraph: currentGraph.metadata!.id!,
         },
       }));
     }
+    // Changes a lot, hopefully okay
   }, [currentGraph]);
 
   // Make sure opened projects sorted ids are in sync with opened projects

packages/core/src/model/EditorDefinition.ts

@@ -158,6 +158,14 @@ export type EditorDefinitionGroup<T extends ChartNode> = SharedEditorDefinitionP
   editors: EditorDefinition<T>[];
 };
 
+export type CustomEditorDefinition<T extends ChartNode> = SharedEditorDefinitionProps<T> & {
+  type: 'custom';
+  customEditorId: string;
+  dataKey?: DataOfType<T, any>;
+  useInputToggleDataKey?: DataOfType<T, boolean>;
+  data?: any;
+};
+
 export type EditorDefinition<T extends ChartNode> =
   | StringEditorDefinition<T>
   | ToggleEditorDefinition<T>
@@ -174,4 +182,5 @@ export type EditorDefinition<T extends ChartNode> =
   | DatasetSelectorEditorDefinition<T>
   | KeyValuePairEditorDefinition<T>
   | EditorDefinitionGroup<T>
-  | StringListEditorDefinition<T>;
+  | StringListEditorDefinition<T>
+  | CustomEditorDefinition<T>;

packages/core/src/model/nodes/ChatNode.ts

@@ -704,7 +704,6 @@ export class ChatNodeImpl extends NodeImpl<ChatNode> {
                 },
               ],
             };
-            console.dir({ output });
           }
 
           const endTime = Date.now();

packages/core/src/model/nodes/CodeNode.ts

@@ -83,6 +83,11 @@ export class CodeNodeImpl extends NodeImpl<CodeNode> {
 
   getEditors(): EditorDefinition<CodeNode>[] {
     return [
+      {
+        type: 'custom',
+        customEditorId: 'CodeNodeAIAssist',
+        label: 'AI Assist',
+      },
       {
         type: 'code',
         label: 'Code',

packages/core/src/model/nodes/GptFunctionNode.ts

@@ -174,7 +174,6 @@ export class GptFunctionNodeImpl extends NodeImpl<GptFunctionNode> {
 
       const interpolated = interpolate(this.data.schema, inputMap);
 
-      console.dir({ inputMap, interpolated, inputs });
       schema = JSON.parse(interpolated);
     }