Make a Position enum.

lynxkite-app/src/lynxkite_app/crdt.py

@@ -200,7 +200,7 @@ def try_to_load_workspace(ws: pycrdt.Map, name: str):
         name: Name of the workspace to load.
     """
     if os.path.exists(name):
-        ws_pyd = workspace.load(name)
+        ws_pyd = workspace.Workspace.load(name)
         crdt_update(
             ws,
             ws_pyd.model_dump(),
@@ -263,18 +263,18 @@ async def execute(name: str, ws_crdt: pycrdt.Map, ws_pyd: workspace.Workspace, d
     path = cwd / name
     assert path.is_relative_to(cwd), "Provided workspace path is invalid"
     # Save user changes before executing, in case the execution fails.
-    workspace.save(ws_pyd, path)
+    ws_pyd.save(path)
     ops.load_user_scripts(name)
-    workspace.connect_crdt(ws_pyd, ws_crdt)
-    workspace.update_metadata(ws_pyd)
-    if not workspace.has_executor(ws_pyd):
+    ws_pyd.connect_crdt(ws_crdt)
+    ws_pyd.update_metadata()
+    if not ws_pyd.has_executor():
         return
     with ws_crdt.doc.transaction():
         for nc in ws_crdt["nodes"]:
             nc["data"]["status"] = "planned"
-    ws_pyd = ws_pyd.normalize()
-    await workspace.execute(ws_pyd)
-    workspace.save(ws_pyd, path)
+    ws_pyd.normalize()
+    await ws_pyd.execute()
+    ws_pyd.save(path)
     print(f"Finished running {name} in {ws_pyd.env}.")
 
 

lynxkite-app/src/lynxkite_app/main.py

@@ -50,13 +50,13 @@ data_path = pathlib.Path()
 def save(req: SaveRequest):
     path = data_path / req.path
     assert path.is_relative_to(data_path)
-    workspace.save(req.ws, path)
+    req.ws.save(path)
 
 
 @app.post("/api/save")
 async def save_and_execute(req: SaveRequest):
     save(req)
-    await workspace.execute(req.ws)
+    await req.ws.execute()
     save(req)
     return req.ws
 
@@ -76,7 +76,7 @@ def load(path: str):
     assert path.is_relative_to(data_path)
     if not path.exists():
         return workspace.Workspace()
-    return workspace.load(path)
+    return workspace.Workspace.load(path)
 
 
 class DirectoryEntry(pydantic.BaseModel):

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

@@ -95,17 +95,28 @@ class ParameterGroup(BaseConfig):
     type: str = "group"
 
 
+class Position(enum.Enum):
+    """Defines the position of an input or output in the UI."""
+
+    LEFT = "left"
+    RIGHT = "right"
+    TOP = "top"
+    BOTTOM = "bottom"
+
+    def is_vertical(self):
+        return self in (self.TOP, self.BOTTOM)
+
+
 class Input(BaseConfig):
     name: str
     type: Type
-    # TODO: Make position an enum with the possible values.
-    position: str = "left"
+    position: Position = Position.LEFT
 
 
 class Output(BaseConfig):
     name: str
     type: Type
-    position: str = "right"
+    position: Position = Position.RIGHT
 
 
 @dataclass
@@ -291,7 +302,7 @@ def input_position(**kwargs):
     def decorator(func):
         op = func.__op__
         for k, v in kwargs.items():
-            op.get_input(k).position = v
+            op.get_input(k).position = Position(v)
         return func
 
     return decorator
@@ -303,7 +314,7 @@ def output_position(**kwargs):
     def decorator(func):
         op = func.__op__
         for k, v in kwargs.items():
-            op.get_output(k).position = v
+            op.get_output(k).position = Position(v)
         return func
 
     return decorator

lynxkite-core/src/lynxkite/core/workspace.py

@@ -110,102 +110,91 @@ class Workspace(BaseConfig):
                     valid_targets.add((n.id, h.name))
                 for h in _ops[n.id].outputs:
                     valid_sources.add((n.id, h.name))
-        edges = [
+        self.edges = [
             edge
             for edge in self.edges
             if (edge.source, edge.sourceHandle) in valid_sources
             and (edge.target, edge.targetHandle) in valid_targets
         ]
-        return self.model_copy(update={"edges": edges})
 
+    def has_executor(self):
+        return self.env in ops.EXECUTORS
+
+    async def execute(self):
+        await ops.EXECUTORS[self.env](self)
+
+    def save(self, path: str):
+        """Persist the workspace to a local file in JSON format."""
+        j = self.model_dump()
+        j = json.dumps(j, indent=2, sort_keys=True) + "\n"
+        dirname, basename = os.path.split(path)
+        if dirname:
+            os.makedirs(dirname, exist_ok=True)
+        # Create temp file in the same directory to make sure it's on the same filesystem.
+        with tempfile.NamedTemporaryFile(
+            "w", encoding="utf-8", prefix=f".{basename}.", dir=dirname, delete=False
+        ) as f:
+            temp_name = f.name
+            f.write(j)
+        os.replace(temp_name, path)
+
+    @staticmethod
+    def load(path: str) -> "Workspace":
+        """Load a workspace from a file.
+
+        After loading the workspace, the metadata of the workspace is updated.
+
+        Args:
+            path (str): The path to the file to load the workspace from.
+
+        Returns:
+            Workspace: The loaded workspace object, with updated metadata.
+        """
+        with open(path, encoding="utf-8") as f:
+            j = f.read()
+        ws = Workspace.model_validate_json(j)
+        # Metadata is added after loading. This way code changes take effect on old boxes too.
+        ws.update_metadata()
+        return ws
 
-def has_executor(ws: Workspace):
-    return ws.env in ops.EXECUTORS
-
-
-async def execute(ws: Workspace):
-    await ops.EXECUTORS[ws.env](ws)
-
-
-def save(ws: Workspace, path: str):
-    """Persist a workspace to a local file in JSON format."""
-    j = ws.model_dump()
-    j = json.dumps(j, indent=2, sort_keys=True) + "\n"
-    dirname, basename = os.path.split(path)
-    if dirname:
-        os.makedirs(dirname, exist_ok=True)
-    # Create temp file in the same directory to make sure it's on the same filesystem.
-    with tempfile.NamedTemporaryFile(
-        "w", encoding="utf-8", prefix=f".{basename}.", dir=dirname, delete=False
-    ) as f:
-        temp_name = f.name
-        f.write(j)
-    os.replace(temp_name, path)
-
-
-def load(path: str) -> Workspace:
-    """Load a workspace from a file.
-
-    After loading the workspace, the metadata of the workspace is updated.
-
-    Args:
-        path (str): The path to the file to load the workspace from.
-
-    Returns:
-        Workspace: The loaded workspace object, with updated metadata.
-    """
-    with open(path, encoding="utf-8") as f:
-        j = f.read()
-    ws = Workspace.model_validate_json(j)
-    # Metadata is added after loading. This way code changes take effect on old boxes too.
-    update_metadata(ws)
-    return ws
-
-
-def update_metadata(ws: Workspace) -> Workspace:
-    """Update the metadata of the given workspace object.
-
-    The metadata is the information about the operations that the nodes in the workspace represent,
-    like the parameters and their possible values.
-    This information comes from the catalog of operations for the environment of the workspace.
-
-    Args:
-        ws: The workspace object to update.
+    def update_metadata(self):
+        """Update the metadata of this workspace.
 
-    Returns:
-        Workspace: The updated workspace object.
-    """
-    if ws.env not in ops.CATALOGS:
-        return ws
-    catalog = ops.CATALOGS[ws.env]
-    for node in ws.nodes:
-        data = node.data
-        op = catalog.get(data.title)
-        if op:
-            if data.meta != op:
-                data.meta = op
-                if hasattr(node, "_crdt"):
-                    node._crdt["data"]["meta"] = op.model_dump()
-            if node.type != op.type:
-                node.type = op.type
-                if hasattr(node, "_crdt"):
-                    node._crdt["type"] = op.type
-            if data.error == "Unknown operation.":
-                data.error = None
+        The metadata is the information about the operations that the nodes in the workspace represent,
+        like the parameters and their possible values.
+        This information comes from the catalog of operations for the environment of the workspace.
+        """
+        if self.env not in ops.CATALOGS:
+            return self
+        catalog = ops.CATALOGS[self.env]
+        for node in self.nodes:
+            data = node.data
+            op = catalog.get(data.title)
+            if op:
+                if data.meta != op:
+                    data.meta = op
+                    if hasattr(node, "_crdt"):
+                        # Go through JSON to simplify the types. CRDT can't handle enums.
+                        node._crdt["data"]["meta"] = json.loads(op.model_dump_json())
+                        print("set metadata to", op)
+                if node.type != op.type:
+                    node.type = op.type
+                    if hasattr(node, "_crdt"):
+                        node._crdt["type"] = op.type
+                if data.error == "Unknown operation.":
+                    data.error = None
+                    if hasattr(node, "_crdt"):
+                        node._crdt["data"]["error"] = None
+            else:
+                data.error = "Unknown operation."
                 if hasattr(node, "_crdt"):
-                    node._crdt["data"]["error"] = None
-        else:
-            data.error = "Unknown operation."
-            if hasattr(node, "_crdt"):
-                node._crdt["data"]["meta"] = {}
-                node._crdt["data"]["error"] = "Unknown operation."
-    return ws
-
-
-def connect_crdt(ws_pyd: Workspace, ws_crdt: pycrdt.Map):
-    ws_pyd._crdt = ws_crdt
-    with ws_crdt.doc.transaction():
-        for nc, np in zip(ws_crdt["nodes"], ws_pyd.nodes):
-            if "data" not in nc:
-                nc["data"] = pycrdt.Map()
-            np._crdt = nc
+                    node._crdt["data"]["meta"] = {}
+                    node._crdt["data"]["error"] = "Unknown operation."
+
+    def connect_crdt(self, ws_crdt: pycrdt.Map):
+        self._crdt = ws_crdt
+        with ws_crdt.doc.transaction():
+            for nc, np in zip(ws_crdt["nodes"], self.nodes):
+                if "data" not in nc:
+                    nc["data"] = pycrdt.Map()
+                np._crdt = nc

lynxkite-core/tests/test_workspace.py

@@ -34,9 +34,9 @@ def test_save_load():
     path = os.path.join(tempfile.gettempdir(), "test_workspace.json")
 
     try:
-        workspace.save(ws, path)
+        ws.save(path)
         assert os.path.exists(path)
-        loaded_ws = workspace.load(path)
+        loaded_ws = workspace.Workspace.load(path)
         assert loaded_ws.env == ws.env
         assert len(loaded_ws.nodes) == len(ws.nodes)
         assert len(loaded_ws.edges) == len(ws.edges)
@@ -88,14 +88,14 @@ def test_update_metadata():
             position=workspace.Position(x=0, y=0),
         )
     )
-    updated_ws = workspace._update_metadata(ws)
-    assert updated_ws.nodes[0].data.meta.name == "Test Operation"
-    assert updated_ws.nodes[0].data.error is None
-    assert not hasattr(updated_ws.nodes[1].data, "meta")
-    assert updated_ws.nodes[1].data.error == "Unknown operation."
+    ws.update_metadata()
+    assert ws.nodes[0].data.meta.name == "Test Operation"
+    assert ws.nodes[0].data.error is None
+    assert not hasattr(ws.nodes[1].data, "meta")
+    assert ws.nodes[1].data.error == "Unknown operation."
 
 
 def test_update_metadata_with_empty_workspace():
     ws = workspace.Workspace(env="test")
-    updated_ws = workspace._update_metadata(ws)
-    assert len(updated_ws.nodes) == 0
+    ws.update_metadata()
+    assert len(ws.nodes) == 0

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

@@ -22,7 +22,7 @@ def load_ws(model_workspace: str):
     path = cwd / model_workspace
     assert path.is_relative_to(cwd)
     assert path.exists(), f"Workspace {path} does not exist"
-    ws = workspace.load(path)
+    ws = workspace.Workspace.load(path)
     return ws
 
 

lynxkite-graph-analytics/src/lynxkite_graph_analytics/pytorch/pytorch_core.py

@@ -21,9 +21,9 @@ def op(name, weights=False, **kwargs):
         _op(func)
         op = func.__op__
         for p in op.inputs:
-            p.position = "bottom"
+            p.position = ops.Position.BOTTOM
         for p in op.outputs:
-            p.position = "top"
+            p.position = ops.Position.TOP
         return func
 
     return decorator

lynxkite-graph-analytics/src/lynxkite_graph_analytics/pytorch/pytorch_ops.py

@@ -151,9 +151,9 @@ ops.register_passive_op(
 def _set_handle_positions(op):
     op: ops.Op = op.__op__
     for v in op.outputs:
-        v.position = "top"
+        v.position = ops.Position.TOP
     for v in op.inputs:
-        v.position = "bottom"
+        v.position = ops.Position.BOTTOM
 
 
 def _register_simple_pytorch_layer(func):

lynxkite-lynxscribe/src/lynxkite_lynxscribe/lynxscribe_ops.py

@@ -875,7 +875,7 @@ async def get_chat_api(ws: str):
     path = cwd / ws
     assert path.is_relative_to(cwd)
     assert path.exists(), f"Workspace {path} does not exist"
-    ws = workspace.load(path)
+    ws = workspace.Workspace.load(path)
     contexts = await ops.EXECUTORS[ENV](ws)
     nodes = [op for op in ws.nodes if op.data.title == "LynxScribe RAG Graph Chatbot Backend"]
     [node] = nodes
@@ -939,7 +939,7 @@ def get_lynxscribe_workspaces():
     for p in pathlib.Path().glob("**/*"):
         if p.is_file():
             try:
-                ws = workspace.load(p)
+                ws = workspace.Workspace.load(p)
                 if ws.env == ENV:
                     workspaces.append(p)
             except Exception: