47 smooth graph creation flow

* Add Graph creation box

* Wrap Operation result into a Result class

* Complex Yjs objects for parameters dont work well, use JSON string representation instead

---------

Co-authored-by: JMLizano <[email protected]>

.gitignore

@@ -15,4 +15,4 @@ build
 joblib-cache
 *.egg-info
 
-lynxkite-app/crdt_data
+lynxkite_crdt_data

lynxkite-app/.gitignore

@@ -1,3 +1,4 @@
-/src/lynxkite/app/web_assets
-!/src/lynxkite/app/web_assets/__init__.py
-!/src/lynxkite/app/web_assets/assets/__init__.py
+/src/lynxkite_app/web_assets
+!/src/lynxkite_app/web_assets/__init__.py
+!/src/lynxkite_app/web_assets/assets/__init__.py
+data/

lynxkite-app/web/src/index.css

@@ -347,3 +347,85 @@ path.react-flow__edge-path {
   outline: var(--xy-selection-border, var(--xy-selection-border-default));
   outline-offset: 7.5px;
 }
+
+.graph-creation-view {
+  display: flex;
+  width: 100%;
+  margin-top: 10px;
+}
+
+.graph-tables, .graph-relations {
+    flex: 1;
+    padding-left: 10px;
+    padding-right: 10px; 
+}
+
+.graph-table-header{
+  display: flex;
+  justify-content: space-between;
+  font-weight: bold;
+  text-align: left;
+  background-color: #333;
+  color: white;
+  padding: 10px;
+  border-bottom: 2px solid #222;
+  font-size: 16px;
+}
+
+.graph-creation-view .df-head {
+  font-weight: bold;
+  display: flex;
+  justify-content: space-between;
+  padding: 8px 12px;
+  border-bottom: 1px solid #ccc; /* Adds a separator between rows */
+}
+
+/* Alternating background colors for table-like effect */
+.graph-creation-view .df-head:nth-child(odd) {
+  background-color: #f9f9f9;
+}
+.graph-creation-view .df-head:nth-child(even) {
+  background-color: #e0e0e0;
+}
+
+.graph-relation-attributes {
+  display: flex;
+  flex-direction: column;
+  gap: 10px; /* Adds space between each label-input pair */
+  width: 100%;
+}
+
+.graph-relation-attributes label {
+  font-size: 12px;
+  font-weight: bold;
+  display: block;
+  margin-bottom: 2px;
+  color: #666;  /* Lighter text for labels */
+}
+
+.graph-relation-attributes input {
+  width: 100%;
+  padding: 8px;
+  font-size: 14px;
+  border: 1px solid #ccc;
+  border-radius: 4px;
+  outline: none;
+}
+
+.graph-relation-attributes input:focus {
+  border-color: #007bff; /* Highlight input on focus */
+}
+
+.add-relationship-button {
+  background-color: #28a745;
+  color: white;
+  border: none;
+  font-size: 16px;
+  cursor: pointer;
+  padding: 4px 10px;
+  border-radius: 4px;
+}
+
+.add-relationship-button:hover {
+  background-color: #218838;
+}

lynxkite-app/web/src/workspace/Workspace.tsx

@@ -45,6 +45,7 @@ import NodeWithImage from "./nodes/NodeWithImage.tsx";
 import NodeWithParams from "./nodes/NodeWithParams";
 import NodeWithTableView from "./nodes/NodeWithTableView.tsx";
 import NodeWithVisualization from "./nodes/NodeWithVisualization.tsx";
+import NodeWithGraphCreationView from "./nodes/GraphCreationNode.tsx";
 
 export default function (props: any) {
   return (
@@ -172,6 +173,7 @@ function LynxKiteFlow() {
       visualization: NodeWithVisualization,
       image: NodeWithImage,
       table_view: NodeWithTableView,
+      graph_creation_view: NodeWithGraphCreationView,
     }),
     [],
   );

lynxkite-app/web/src/workspace/nodes/GraphCreationNode.tsx

@@ -0,0 +1,308 @@
+import { useReactFlow } from "@xyflow/react";
+import { useState } from "react";
+import React from "react";
+import Markdown from "react-markdown";
+// @ts-ignore
+import Trash from "~icons/tabler/trash";
+import LynxKiteNode from "./LynxKiteNode";
+import Table from "./Table";
+
+function toMD(v: any): string {
+  if (typeof v === "string") {
+    return v;
+  }
+  if (Array.isArray(v)) {
+    return v.map(toMD).join("\n\n");
+  }
+  return JSON.stringify(v);
+}
+
+function displayTable(name: string, df: any) {
+  if (df.data.length > 1) {
+    return (
+      <Table
+        key={`${name}-table`}
+        name={`${name}-table`}
+        columns={df.columns}
+        data={df.data}
+      />
+    );
+  }
+  if (df.data.length) {
+    return (
+      <dl key={`${name}-dl`}>
+        {df.columns.map((c: string, i: number) => (
+          <React.Fragment key={`${name}-${c}`}>
+            <dt>{c}</dt>
+            <dd>
+              <Markdown>{toMD(df.data[0][i])}</Markdown>
+            </dd>
+          </React.Fragment>
+        ))}
+      </dl>
+    );
+  }
+  return JSON.stringify(df.data);
+}
+
+function relationsToDict(relations: any[]) {
+  if (!relations) {
+    return {};
+  }
+  return Object.assign({}, ...relations.map((r: any) => ({ [r.name]: r })));
+}
+
+export type UpdateOptions = { delay?: number };
+
+export default function NodeWithGraphCreationView(props: any) {
+  const reactFlow = useReactFlow();
+  const [open, setOpen] = useState({} as { [name: string]: boolean });
+  const display = props.data.display?.value;
+  const tables = display?.dataframes || {};
+  const singleTable = tables && Object.keys(tables).length === 1;
+  const [relations, setRelations] = useState(
+    relationsToDict(display?.relations) || {},
+  );
+  const singleRelation = relations && Object.keys(relations).length === 1;
+  function setParam(name: string, newValue: any, opts: UpdateOptions) {
+    reactFlow.updateNodeData(props.id, {
+      params: { ...props.data.params, [name]: newValue },
+      __execution_delay: opts.delay || 0,
+    });
+  }
+
+  function updateRelation(event: any, relation: any) {
+    event.preventDefault();
+
+    const updatedRelation = {
+      ...relation,
+      ...Object.fromEntries(new FormData(event.target).entries()),
+    };
+
+    // Avoid mutating React state directly
+    const newRelations = { ...relations };
+    if (relation.name !== updatedRelation.name) {
+      delete newRelations[relation.name];
+    }
+    newRelations[updatedRelation.name] = updatedRelation;
+    setRelations(newRelations);
+    // There is some issue with how Yjs handles complex objects (maps, arrays)
+    // so we need to serialize the relations object to a string
+    setParam("relations", JSON.stringify(newRelations), {});
+  }
+
+  const addRelation = () => {
+    const new_relation = {
+      name: "new_relation",
+      df: "",
+      source_column: "",
+      target_column: "",
+      source_table: "",
+      target_table: "",
+      source_key: "",
+      target_key: "",
+    };
+    setRelations({
+      ...relations,
+      [new_relation.name]: new_relation,
+    });
+    setOpen({ ...open, [new_relation.name]: true });
+  };
+
+  const deleteRelation = (relation: any) => {
+    const newOpen = { ...open };
+    delete newOpen[relation.name];
+    setOpen(newOpen);
+    const newRelations = { ...relations };
+    delete newRelations[relation.name];
+    setRelations(newRelations);
+    // There is some issue with how Yjs handles complex objects (maps, arrays)
+    // so we need to serialize the relations object to a string
+    setParam("relations", JSON.stringify(newRelations), {});
+  };
+
+  function displayRelation(relation: any) {
+    // TODO: Dynamic autocomplete
+    return (
+      <form
+        className="graph-relation-attributes"
+        onSubmit={(e) => {
+          updateRelation(e, relation);
+        }}
+      >
+        <label htmlFor="name">Name:</label>
+        <input type="text" id="name" name="name" defaultValue={relation.name} />
+
+        <label htmlFor="df">DataFrame:</label>
+        <input
+          type="text"
+          id="df"
+          name="df"
+          defaultValue={relation.df}
+          list="df-options"
+          required
+        />
+
+        <label htmlFor="source_column">Source Column:</label>
+        <input
+          type="text"
+          id="source_column"
+          name="source_column"
+          defaultValue={relation.source_column}
+          list="edges-column-options"
+          required
+        />
+
+        <label htmlFor="target_column">Target Column:</label>
+        <input
+          type="text"
+          id="target_column"
+          name="target_column"
+          defaultValue={relation.target_column}
+          list="edges-column-options"
+          required
+        />
+
+        <label htmlFor="source_table">Source Table:</label>
+        <input
+          type="text"
+          id="source_table"
+          name="source_table"
+          defaultValue={relation.source_table}
+          list="df-options"
+          required
+        />
+
+        <label htmlFor="target_table">Target Table:</label>
+        <input
+          type="text"
+          id="target_table"
+          name="target_table"
+          defaultValue={relation.target_table}
+          list="df-options"
+          required
+        />
+
+        <label htmlFor="source_key">Source Key:</label>
+        <input
+          type="text"
+          id="source_key"
+          name="source_key"
+          defaultValue={relation.source_key}
+          list="source-node-column-options"
+          required
+        />
+
+        <label htmlFor="target_key">Target Key:</label>
+        <input
+          type="text"
+          id="target_key"
+          name="target_key"
+          defaultValue={relation.target_key}
+          list="target-node-column-options"
+          required
+        />
+
+        <datalist id="df-options">
+          {Object.keys(tables).map((name) => (
+            <option key={name} value={name} />
+          ))}
+        </datalist>
+
+        <datalist id="edges-column-options">
+          {tables[relation.source_table] &&
+            tables[relation.df].columns.map((name: string) => (
+              <option key={name} value={name} />
+            ))}
+        </datalist>
+
+        <datalist id="source-node-column-options">
+          {tables[relation.source_table] &&
+            tables[relation.source_table].columns.map((name: string) => (
+              <option key={name} value={name} />
+            ))}
+        </datalist>
+
+        <datalist id="target-node-column-options">
+          {tables[relation.source_table] &&
+            tables[relation.target_table].columns.map((name: string) => (
+              <option key={name} value={name} />
+            ))}
+        </datalist>
+
+        <button className="submit-relationship-button" type="submit">
+          Create
+        </button>
+      </form>
+    );
+  }
+
+  return (
+    <LynxKiteNode {...props}>
+      <div className="graph-creation-view">
+        <div className="graph-tables">
+          <div className="graph-table-header">Node Tables</div>
+          {display && [
+            Object.entries(tables).map(([name, df]: [string, any]) => (
+              <React.Fragment key={name}>
+                {!singleTable && (
+                  <div
+                    key={`${name}-header`}
+                    className="df-head"
+                    onClick={() => setOpen({ ...open, [name]: !open[name] })}
+                  >
+                    {name}
+                  </div>
+                )}
+                {(singleTable || open[name]) && displayTable(name, df)}
+              </React.Fragment>
+            )),
+            Object.entries(display.others || {}).map(([name, o]) => (
+              <>
+                <div
+                  key={name}
+                  className="df-head"
+                  onClick={() => setOpen({ ...open, [name]: !open[name] })}
+                >
+                  {name}
+                </div>
+                {open[name] && <pre>{(o as any).toString()}</pre>}
+              </>
+            )),
+          ]}
+        </div>
+        <div className="graph-relations">
+          <div className="graph-table-header">
+            Relationships
+            <button
+              className="add-relationship-button"
+              onClick={(_) => addRelation()}
+            >
+              +
+            </button>
+          </div>
+          {relations &&
+            Object.entries(relations).map(([name, relation]: [string, any]) => (
+              <React.Fragment key={name}>
+                <div
+                  key={`${name}-header`}
+                  className="df-head"
+                  onClick={() => setOpen({ ...open, [name]: !open[name] })}
+                >
+                  {name}
+                  <button
+                    onClick={() => {
+                      deleteRelation(relation);
+                    }}
+                  >
+                    <Trash />
+                  </button>
+                </div>
+                {(singleRelation || open[name]) && displayRelation(relation)}
+              </React.Fragment>
+            ))}
+        </div>
+      </div>
+    </LynxKiteNode>
+  );
+}

lynxkite-app/web/src/workspace/nodes/Table.tsx

@@ -1,6 +1,6 @@
 export default function Table(props: any) {
   return (
-    <table>
+    <table id={props.name || "table"}>
       <thead>
         <tr>
           {props.columns.map((column: string) => (

lynxkite-app/web/tests/directory.spec.ts

@@ -1,4 +1,4 @@
-// Tests some basic operations like box creation, deletion, and dragging.
+// Tests the basic directory operations, such as creating and deleting folders and workspaces.
 import { expect, test } from "@playwright/test";
 import { Splash, Workspace } from "./lynxkite";
 

lynxkite-app/web/tests/graph_creation.spec.ts

@@ -0,0 +1,89 @@
+// Test the graph creation box in LynxKite
+import { expect, test } from "@playwright/test";
+import { Splash, Workspace } from "./lynxkite";
+
+let workspace: Workspace;
+
+test.beforeEach(async ({ browser }) => {
+  workspace = await Workspace.empty(
+    await browser.newPage(),
+    "graph_creation_spec_test",
+  );
+  await workspace.addBox("Create scale-free graph");
+  await workspace.addBox("Create graph");
+  await workspace.connectBoxes("Create scale-free graph 1", "Create graph 1");
+});
+
+test.afterEach(async () => {
+  await workspace.close();
+  const splash = await new Splash(workspace.page);
+  splash.page.on("dialog", async (dialog) => {
+    await dialog.accept();
+  });
+  await splash.deleteEntry("graph_creation_spec_test");
+});
+
+test("Tables are displayed in the Graph creation box", async () => {
+  const graphBox = await workspace.getBox("Create graph 1");
+  const nodesTableHeader = await graphBox.locator(".graph-tables .df-head", {
+    hasText: "nodes",
+  });
+  const edgesTableHeader = await graphBox.locator(".graph-tables .df-head", {
+    hasText: "edges",
+  });
+  await expect(nodesTableHeader).toBeVisible();
+  await expect(edgesTableHeader).toBeVisible();
+  nodesTableHeader.click();
+  await expect(graphBox.locator("#nodes-table")).toBeVisible();
+  edgesTableHeader.click();
+  await expect(graphBox.locator("#edges-table")).toBeVisible();
+});
+
+test("Adding and removing relationships", async () => {
+  const graphBox = await workspace.getBox("Create graph 1");
+  const addRelationshipButton = await graphBox.locator(
+    ".add-relationship-button",
+  );
+  await addRelationshipButton.click();
+  const formData: Record<string, string> = {
+    name: "relation_1",
+    df: "edges",
+    source_column: "source_id",
+    target_column: "target_id",
+    source_table: "nodes",
+    target_table: "nodes",
+    source_key: "node_id",
+    target_key: "node_id",
+  };
+  for (const [fieldName, fieldValue] of Object.entries(formData)) {
+    const inputField = await graphBox.locator(
+      `.graph-relation-attributes input[name="${fieldName}"]`,
+    );
+    await inputField.fill(fieldValue);
+  }
+  await graphBox.locator(".submit-relationship-button").click();
+  // check that the relationship has been saved in the backend
+  await workspace.page.reload();
+  const graphBoxAfterReload = await workspace.getBox("Create graph 1");
+  const relationHeader = await graphBoxAfterReload.locator(
+    ".graph-relations .df-head",
+    { hasText: "relation_1" },
+  );
+  await expect(relationHeader).toBeVisible();
+  await relationHeader.locator("button").click(); // Delete the relationship
+  await expect(relationHeader).not.toBeVisible();
+});
+
+test("Output of the box is a bundle", async () => {
+  await workspace.addBox("View tables");
+  const tableView = await workspace.getBox("View tables 1");
+  await workspace.connectBoxes("Create graph 1", "View tables 1");
+  const nodesTableHeader = await tableView.locator(".df-head", {
+    hasText: "nodes",
+  });
+  const edgesTableHeader = await tableView.locator(".df-head", {
+    hasText: "edges",
+  });
+  await expect(nodesTableHeader).toBeVisible();
+  await expect(edgesTableHeader).toBeVisible();
+});

lynxkite-app/web/tests/lynxkite.ts

@@ -57,8 +57,9 @@ export class Workspace {
     const allBoxes = await this.getBoxes();
     if (allBoxes) {
       // Avoid overlapping with existing nodes
-      const numNodes = allBoxes.length;
-      await this.page.mouse.wheel(0, numNodes * 500);
+      const numNodes = allBoxes.length || 1;
+      await this.page.mouse.wheel(0, numNodes * 400);
+      await new Promise((resolve) => setTimeout(resolve, 200));
     }
 
     // Some x,y offset, otherwise the box handle may fall outside the viewport.
@@ -97,7 +98,12 @@ export class Workspace {
     return this.page.locator(".react-flow__node").all();
   }
 
-  getBoxHandle(boxId: string) {
+  getBoxHandle(boxId: string, pos?: string) {
+    if (pos) {
+      return this.page.locator(
+        `[data-id="${boxId}"] [data-handlepos="${pos}"]`,
+      );
+    }
     return this.page.getByTestId(boxId);
   }
 
@@ -130,8 +136,8 @@ export class Workspace {
   }
 
   async connectBoxes(sourceId: string, targetId: string) {
-    const sourceHandle = this.getBoxHandle(sourceId);
-    const targetHandle = this.getBoxHandle(targetId);
+    const sourceHandle = this.getBoxHandle(sourceId, "right");
+    const targetHandle = this.getBoxHandle(targetId, "left");
     await sourceHandle.hover();
     await this.page.mouse.down();
     await targetHandle.hover();

lynxkite-core/src/lynxkite/core/executors/one_by_one.py

@@ -141,28 +141,27 @@ async def execute(ws: workspace.Workspace, catalog, cache=None):
                     if cache is not None:
                         key = make_cache_key((inputs, params))
                         if key not in cache:
-                            cache[key] = await await_if_needed(op(*inputs, **params))
-                        result = cache[key]
+                            result: ops.Result = op(*inputs, **params)
+                            output = await await_if_needed(result.output)
+                            cache[key] = output
+                        output = cache[key]
                     else:
-                        result = await await_if_needed(op(*inputs, **params))
+                        result = op(*inputs, **params)
+                        output = await await_if_needed(result.output)
                 except Exception as e:
                     traceback.print_exc()
                     data.error = str(e)
                     break
-                contexts[node.id].last_result = result
+                contexts[node.id].last_result = output
                 # Returned lists and DataFrames are considered multiple tasks.
-                if isinstance(result, pd.DataFrame):
-                    result = df_to_list(result)
-                elif not isinstance(result, list):
-                    result = [result]
-                results.extend(result)
+                if isinstance(output, pd.DataFrame):
+                    output = df_to_list(output)
+                elif not isinstance(output, list):
+                    output = [output]
+                results.extend(output)
             else:  # Finished all tasks without errors.
-                if (
-                    op.type == "visualization"
-                    or op.type == "table_view"
-                    or op.type == "image"
-                ):
-                    data.display = results[0]
+                if result.display:
+                    data.display = await await_if_needed(result.display)
                 for edge in edges[node.id]:
                     t = nodes[edge.target]
                     op = catalog[t.data.title]

lynxkite-core/src/lynxkite/core/ops.py

@@ -6,6 +6,7 @@ import functools
 import inspect
 import pydantic
 import typing
+from dataclasses import dataclass
 from typing_extensions import Annotated
 
 CATALOGS = {}
@@ -28,6 +29,16 @@ PathStr = Annotated[str, {"format": "path"}]
 CollapsedStr = Annotated[str, {"format": "collapsed"}]
 NodeAttribute = Annotated[str, {"format": "node attribute"}]
 EdgeAttribute = Annotated[str, {"format": "edge attribute"}]
+# https://github.com/python/typing/issues/182#issuecomment-1320974824
+ReadOnlyJSON: typing.TypeAlias = (
+    typing.Mapping[str, "ReadOnlyJSON"]
+    | typing.Sequence["ReadOnlyJSON"]
+    | str
+    | int
+    | float
+    | bool
+    | None
+)
 
 
 class BaseConfig(pydantic.BaseModel):
@@ -74,6 +85,19 @@ class Output(BaseConfig):
     position: str = "right"
 
 
+@dataclass
+class Result:
+    """Represents the result of an operation.
+
+    The `output` attribute is what will be used as input for other operations.
+    The `display` attribute is used to send data to display on the UI. The value has to be
+    JSON-serializable.
+    """
+
+    output: typing.Any
+    display: ReadOnlyJSON | None = None
+
+
 MULTI_INPUT = Input(name="multi", type="*")
 
 
@@ -105,6 +129,18 @@ class Op(BaseConfig):
                 elif isinstance(self.params[p].type, enum.EnumMeta):
                     params[p] = self.params[p].type[params[p]]
         res = self.func(*inputs, **params)
+        if not isinstance(res, Result):
+            # Automatically wrap the result in a Result object, if it isn't already.
+            res = Result(output=res)
+            if self.type in [
+                "visualization",
+                "table_view",
+                "graph_creation_view",
+                "image",
+            ]:
+                # If the operation is some kind of visualization, we use the output as the
+                # value to display by default.
+                res.display = res.output
         return res
 
 

lynxkite-core/tests/test_ops.py

@@ -76,10 +76,44 @@ def test_op_decorator_with_complex_types():
     assert complex_op.__op__.inputs == {
         "color": ops.Input(name="color", type=Color, position="left"),
         "color_list": ops.Input(name="color_list", type=list[Color], position="left"),
-        "color_dict": ops.Input(name="color_dict", type=dict[str, Color], position="left"),
+        "color_dict": ops.Input(
+            name="color_dict", type=dict[str, Color], position="left"
+        ),
     }
     assert complex_op.__op__.type == "basic"
     assert complex_op.__op__.outputs == {
         "result": ops.Output(name="result", type=None, position="right")
     }
     assert ops.CATALOGS["test"]["color_op"] == complex_op.__op__
+
+
+def test_operation_can_return_non_result_instance():
+    @ops.op(env="test", name="subtract", view="basic", outputs=["result"])
+    def subtract(a, b):
+        return a - b
+
+    result = ops.CATALOGS["test"]["subtract"](5, 3)
+    assert isinstance(result, ops.Result)
+    assert result.output == 2
+    assert result.display is None
+
+
+def test_operation_can_return_result_instance():
+    @ops.op(env="test", name="subtract", view="basic", outputs=["result"])
+    def subtract(a, b):
+        return ops.Result(output=a - b, display=None)
+
+    result = ops.CATALOGS["test"]["subtract"](5, 3)
+    assert isinstance(result, ops.Result)
+    assert result.output == 2
+    assert result.display is None
+
+
+def test_visualization_operations_display_is_populated_automatically():
+    @ops.op(env="test", name="display_op", view="visualization", outputs=["result"])
+    def display_op():
+        return {"display_value": 1}
+
+    result = ops.CATALOGS["test"]["display_op"]()
+    assert isinstance(result, ops.Result)
+    assert result.output == result.display == {"display_value": 1}

lynxkite-graph-analytics/src/lynxkite_graph_analytics/lynxkite_ops.py

@@ -14,6 +14,8 @@ import pandas as pd
 import polars as pl
 import traceback
 import typing
+import json
+
 
 mem = joblib.Memory("../joblib-cache")
 ENV = "LynxKite Graph Analytics"
@@ -35,6 +37,7 @@ class RelationDefinition:
     target_table: str  # The DataFrame that contains the target nodes.
     source_key: str  # The column in the source table that contains the node ID.
     target_key: str  # The column in the target table that contains the node ID.
+    name: str | None = None  # Descriptive name for the relation.
 
 
 @dataclasses.dataclass
@@ -96,6 +99,19 @@ class Bundle:
             other=dict(self.other) if self.other else None,
         )
 
+    def to_dict(self, limit: int = 100):
+        return {
+            "dataframes": {
+                name: {
+                    "columns": [str(c) for c in df.columns],
+                    "data": collect(df)[:limit],
+                }
+                for name, df in self.dfs.items()
+            },
+            "relations": [dataclasses.asdict(relation) for relation in self.relations],
+            "other": self.other,
+        }
+
 
 def nx_node_attribute_func(name):
     """Decorator for wrapping a function that adds a NetworkX node attribute."""
@@ -153,7 +169,7 @@ async def execute(ws):
                             inputs[i] = Bundle.from_nx(x)
                         elif p.type == Bundle and isinstance(x, pd.DataFrame):
                             inputs[i] = Bundle.from_df(x)
-                    output = op(*inputs, **params)
+                    result = op(*inputs, **params)
                 except Exception as e:
                     traceback.print_exc()
                     data.error = str(e)
@@ -163,13 +179,9 @@ async def execute(ws):
                     # It's a flexible input. Create n+1 handles.
                     data.inputs = {f"input{i}": None for i in range(len(inputs) + 1)}
                 data.error = None
-                outputs[node.id] = output
-                if (
-                    op.type == "visualization"
-                    or op.type == "table_view"
-                    or op.type == "image"
-                ):
-                    data.display = output
+                outputs[node.id] = result.output
+                if result.display:
+                    data.display = result.display
 
 
 @op("Import Parquet")
@@ -404,15 +416,33 @@ def collect(df: pd.DataFrame):
 
 @op("View tables", view="table_view")
 def view_tables(bundle: Bundle, *, limit: int = 100):
-    v = {
-        "dataframes": {
-            name: {
-                "columns": [str(c) for c in df.columns],
-                "data": collect(df)[:limit],
-            }
-            for name, df in bundle.dfs.items()
-        },
-        "relations": bundle.relations,
-        "other": bundle.other,
-    }
-    return v
+    return bundle.to_dict(limit=limit)
+
+
+@op(
+    "Create graph",
+    view="graph_creation_view",
+    outputs=["output"],
+)
+def create_graph(bundle: Bundle, *, relations: str = None) -> Bundle:
+    """Replace relations of the given bundle
+
+    relations is a stringified JSON, instead of a dict, because complex Yjs types (arrays, maps)
+    are not currently supported in the UI.
+
+    Args:
+        bundle: Bundle to modify
+        relations (str, optional): Set of relations to set for the bundle. The parameter
+            should be a JSON object where the keys are relation names and the values are
+            a dictionary representation of a `RelationDefinition`.
+        Defaults to None.
+
+    Returns:
+        Bundle: The input bundle with the new relations set.
+    """
+    bundle = bundle.copy()
+    if not (relations is None or relations.strip() == ""):
+        bundle.relations = [
+            RelationDefinition(**r) for r in json.loads(relations).values()
+        ]
+    return ops.Result(output=bundle, display=bundle.to_dict(limit=100))