Adding tests (#50)

* Added test for lynxkite-core

* Added tests for lynxkite-app

* Add tests for lynxkite-graph-analytics

* Store the actual Enum value when creating the CRDT object

* Better error message when workspace path is outside data dir

* Use enum for operation Inputs & Outputs position values instead of hardcoded strings everywhere

* Make WorkspaceNode positon an Enum

* Make Operation type an Enum, and rename it to view_type

* Remove register_area operation, it is not used anymore

* Add some missing docstrings and types

* update apiTypes.ts

---------

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

lynxkite-app/src/lynxkite/app/crdt.py

@@ -29,7 +29,11 @@ def ws_exception_handler(exception, log):
 
 
 class WebsocketServer(pycrdt_websocket.WebsocketServer):
-    async def init_room(self, name):
+    async def init_room(self, name: str) -> pycrdt_websocket.YRoom:
+        """Initialize a room for the workspace with the given name.
+
+        The workspace is loaded from "crdt_data" if it exists there, or from "data", or a new workspace is created.
+        """
         path = CRDT_PATH / f"{name}.crdt"
         assert path.is_relative_to(CRDT_PATH)
         ystore = pycrdt_websocket.ystore.FileYStore(path)
@@ -49,6 +53,8 @@ class WebsocketServer(pycrdt_websocket.WebsocketServer):
             ws["edges"] = pycrdt.Array()
         if "env" not in ws:
             ws["env"] = "unset"
+            # We have two possible sources of truth for the workspaces, the YStore and the JSON files.
+            # In case we didn't find the workspace in the YStore, we try to load it from the JSON files.
             try_to_load_workspace(ws, name)
         room = pycrdt_websocket.YRoom(
             ystore=ystore, ydoc=ydoc, exception_handler=ws_exception_handler
@@ -62,6 +68,12 @@ class WebsocketServer(pycrdt_websocket.WebsocketServer):
         return room
 
     async def get_room(self, name: str) -> pycrdt_websocket.YRoom:
+        """Get a room by name.
+
+        This method overrides the parent get_room method. The original creates an empty room,
+        with no associated Ydoc. Instead, we want to initialize the the room with a Workspace
+        object.
+        """
         if name not in self.rooms:
             self.rooms[name] = await self.init_room(name)
         room = self.rooms[name]
@@ -72,7 +84,7 @@ class WebsocketServer(pycrdt_websocket.WebsocketServer):
 last_ws_input = None
 
 
-def clean_input(ws_pyd):
+def clean_input(ws_pyd: workspace.Workspace):
     for node in ws_pyd.nodes:
         node.data.display = None
         node.data.error = None
@@ -83,21 +95,43 @@ def clean_input(ws_pyd):
                 delattr(node, key)
 
 
-def crdt_update(crdt_obj, python_obj, boxes=set()):
+def crdt_update(
+    crdt_obj: pycrdt.Map | pycrdt.Array,
+    python_obj: dict | list,
+    non_collaborative_fields: set[str] = set(),
+):
+    """Update a CRDT object to match a Python object.
+
+    The types between the CRDT object and the Python object must match. If the Python object
+    is a dict, the CRDT object must be a Map. If the Python object is a list, the CRDT object
+    must be an Array.
+
+    Args:
+        crdt_obj: The CRDT object, that will be updated to match the Python object.
+        python_obj: The Python object to update with.
+        non_collaborative_fields: List of fields to treat as a black box. Black boxes are
+        updated as a whole, instead of having a fine-grained data structure to edit
+        collaboratively. Useful for complex fields that contain auto-generated data or
+        metadata.
+        The default is an empty set.
+
+    Raises:
+        ValueError: If the Python object provided is not a dict or list.
+    """
     if isinstance(python_obj, dict):
         for key, value in python_obj.items():
-            if key in boxes:
+            if key in non_collaborative_fields:
                 crdt_obj[key] = value
             elif isinstance(value, dict):
                 if crdt_obj.get(key) is None:
                     crdt_obj[key] = pycrdt.Map()
-                crdt_update(crdt_obj[key], value, boxes)
+                crdt_update(crdt_obj[key], value, non_collaborative_fields)
             elif isinstance(value, list):
                 if crdt_obj.get(key) is None:
                     crdt_obj[key] = pycrdt.Array()
-                crdt_update(crdt_obj[key], value, boxes)
+                crdt_update(crdt_obj[key], value, non_collaborative_fields)
             elif isinstance(value, enum.Enum):
-                crdt_obj[key] = str(value)
+                crdt_obj[key] = str(value.value)
             else:
                 crdt_obj[key] = value
     elif isinstance(python_obj, list):
@@ -105,12 +139,14 @@ def crdt_update(crdt_obj, python_obj, boxes=set()):
             if isinstance(value, dict):
                 if i >= len(crdt_obj):
                     crdt_obj.append(pycrdt.Map())
-                crdt_update(crdt_obj[i], value, boxes)
+                crdt_update(crdt_obj[i], value, non_collaborative_fields)
             elif isinstance(value, list):
                 if i >= len(crdt_obj):
                     crdt_obj.append(pycrdt.Array())
-                crdt_update(crdt_obj[i], value, boxes)
+                crdt_update(crdt_obj[i], value, non_collaborative_fields)
             else:
+                if isinstance(value, enum.Enum):
+                    value = str(value.value)
                 if i >= len(crdt_obj):
                     crdt_obj.append(value)
                 else:
@@ -119,18 +155,34 @@ def crdt_update(crdt_obj, python_obj, boxes=set()):
         raise ValueError("Invalid type:", python_obj)
 
 
-def try_to_load_workspace(ws, name):
+def try_to_load_workspace(ws: pycrdt.Map, name: str):
+    """Load the workspace `name`, if it exists, and update the `ws` CRDT object to match its contents.
+
+    Args:
+        ws: CRDT object to udpate with the workspace contents.
+        name: Name of the workspace to load.
+    """
     json_path = f"data/{name}"
     if os.path.exists(json_path):
         ws_pyd = workspace.load(json_path)
-        crdt_update(ws, ws_pyd.model_dump(), boxes={"display"})
+        # We treat the display field as a black box, since it is a large
+        # dictionary that is meant to change as a whole.
+        crdt_update(ws, ws_pyd.model_dump(), non_collaborative_fields={"display"})
 
 
 last_known_versions = {}
 delayed_executions = {}
 
 
-async def workspace_changed(name, changes, ws_crdt):
+async def workspace_changed(name: str, changes: pycrdt.MapEvent, ws_crdt: pycrdt.Map):
+    """Callback to react to changes in the workspace.
+
+
+    Args:
+        name: Name of the workspace.
+        changes: Changes performed to the workspace.
+        ws_crdt: CRDT object representing the workspace.
+    """
     ws_pyd = workspace.Workspace.model_validate(ws_crdt.to_py())
     # Do not trigger execution for superficial changes.
     # This is a quick solution until we build proper caching.
@@ -154,22 +206,35 @@ async def workspace_changed(name, changes, ws_crdt):
         await execute(name, ws_crdt, ws_pyd)
 
 
-async def execute(name, ws_crdt, ws_pyd, delay=0):
+async def execute(
+    name: str, ws_crdt: pycrdt.Map, ws_pyd: workspace.Workspace, delay: int = 0
+):
+    """Execute the workspace and update the CRDT object with the results.
+
+    Args:
+        name: Name of the workspace.
+        ws_crdt: CRDT object representing the workspace.
+        ws_pyd: Workspace object to execute.
+        delay: Wait time before executing the workspace. The default is 0.
+    """
     if delay:
         try:
             await asyncio.sleep(delay)
         except asyncio.CancelledError:
             return
     path = DATA_PATH / name
-    assert path.is_relative_to(DATA_PATH)
+    assert path.is_relative_to(DATA_PATH), "Provided workspace path is invalid"
+    # Save user changes before executing, in case the execution fails.
     workspace.save(ws_pyd, path)
     await workspace.execute(ws_pyd)
     workspace.save(ws_pyd, path)
+    # Execution happened on the Python object, we need to replicate
+    # the results to the CRDT object.
     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()
-            # Display is added as an opaque Box.
+            # Display is added as a non collaborative field.
             nc["data"]["display"] = np.data.display
             nc["data"]["error"] = np.data.error
 

lynxkite-app/tests/test_crdt.py

@@ -0,0 +1,72 @@
+from enum import Enum
+import pycrdt
+import pytest
+from lynxkite.app.crdt import crdt_update
+
+
+@pytest.fixture
+def empty_dict_workspace():
+    ydoc = pycrdt.Doc()
+    ydoc["workspace"] = ws = pycrdt.Map()
+    yield ws
+
+
+@pytest.fixture
+def empty_list_workspace():
+    ydoc = pycrdt.Doc()
+    ydoc["workspace"] = ws = pycrdt.Array()
+    yield ws
+
+
+class MyEnum(Enum):
+    VALUE = 1
+
+
+@pytest.mark.parametrize(
+    "python_obj,expected",
+    [
+        (
+            {
+                "key1": "value1",
+                "key2": {
+                    "nested_key1": "nested_value1",
+                    "nested_key2": ["nested_value2"],
+                    "nested_key3": MyEnum.VALUE,
+                },
+            },
+            {
+                "key1": "value1",
+                "key2": {
+                    "nested_key1": "nested_value1",
+                    "nested_key2": ["nested_value2"],
+                    "nested_key3": "1",
+                },
+            },
+        )
+    ],
+)
+def test_crdt_update_with_dict(empty_dict_workspace, python_obj, expected):
+    crdt_update(empty_dict_workspace, python_obj)
+    assert empty_dict_workspace.to_py() == expected
+
+
+@pytest.mark.parametrize(
+    "python_obj,expected",
+    [
+        (
+            [
+                "value1",
+                {"nested_key1": "nested_value1", "nested_key2": ["nested_value2"]},
+                MyEnum.VALUE,
+            ],
+            [
+                "value1",
+                {"nested_key1": "nested_value1", "nested_key2": ["nested_value2"]},
+                "1",
+            ],
+        ),
+    ],
+)
+def test_crdt_update_with_list(empty_list_workspace, python_obj, expected):
+    crdt_update(empty_list_workspace, python_obj)
+    assert empty_list_workspace.to_py() == expected

lynxkite-app/tests/test_main.py

@@ -0,0 +1,77 @@
+import uuid
+from fastapi.testclient import TestClient
+from lynxkite.app.main import app, detect_plugins, DATA_PATH
+import os
+
+
+client = TestClient(app)
+
+
+def test_detect_plugins_with_plugins():
+    # This test assumes that these plugins are installed as part of the testing process.
+    plugins = detect_plugins()
+    assert all(
+        plugin in plugins.keys()
+        for plugin in [
+            "lynxkite_plugins.graph_analytics",
+            "lynxkite_plugins.lynxscribe",
+            "lynxkite_plugins.pillow_example",
+        ]
+    )
+
+
+def test_get_catalog():
+    response = client.get("/api/catalog")
+    assert response.status_code == 200
+
+
+def test_save_and_load():
+    save_request = {
+        "path": "test",
+        "ws": {
+            "env": "test",
+            "nodes": [
+                {
+                    "id": "Node_1",
+                    "type": "basic",
+                    "data": {
+                        "display": None,
+                        "error": "Unknown operation.",
+                        "title": "Test node",
+                        "params": {"param1": "value"},
+                    },
+                    "position": {"x": -493.5496596237119, "y": 20.90123252513356},
+                }
+            ],
+            "edges": [],
+        },
+    }
+    response = client.post("/api/save", json=save_request)
+    saved_ws = response.json()
+    assert response.status_code == 200
+    response = client.get("/api/load?path=test")
+    assert response.status_code == 200
+    assert saved_ws == response.json()
+
+
+def test_list_dir():
+    test_dir = str(uuid.uuid4())
+    test_dir_full_path = DATA_PATH / test_dir
+    test_dir_full_path.mkdir(exist_ok=True)
+    test_file = test_dir_full_path / "test_file.txt"
+    test_file.touch()
+    response = client.get(f"/api/dir/list?path={str(test_dir)}")
+    assert response.status_code == 200
+    assert len(response.json()) == 1
+    assert response.json()[0]["name"] == f"{test_dir}/test_file.txt"
+    assert response.json()[0]["type"] == "workspace"
+    test_file.unlink()
+    test_dir_full_path.rmdir()
+
+
+def test_make_dir():
+    dir_name = str(uuid.uuid4())
+    response = client.post("/api/dir/mkdir", json={"path": dir_name})
+    assert response.status_code == 200
+    assert os.path.exists(DATA_PATH / dir_name)
+    os.rmdir(DATA_PATH / dir_name)

lynxkite-app/web/src/apiTypes.ts

@@ -5,21 +5,21 @@
 /* Do not modify it by hand - just update the pydantic models and then re-run the script
 */
 
-/* eslint-disable */
-/**
- * This file was automatically generated by json-schema-to-typescript.
- * DO NOT MODIFY IT BY HAND. Instead, modify the source JSONSchema file,
- * and run json-schema-to-typescript to regenerate this file.
- */
-
 export interface BaseConfig {
   [k: string]: unknown;
 }
-export interface Position {
+export interface NodePosition {
   x: number;
   y: number;
   [k: string]: unknown;
 }
+/**
+ * A workspace is a representation of a computational graph that consists of nodes and edges.
+ *
+ * Each node represents an operation or task, and the edges represent the flow of data between
+ * the nodes. Each workspace is associated with an environment, which determines the operations
+ * that can be performed in the workspace and the execution method for the operations.
+ */
 export interface Workspace {
   env?: string;
   nodes?: WorkspaceNode[];
@@ -30,7 +30,7 @@ export interface WorkspaceNode {
   id: string;
   type: string;
   data: WorkspaceNodeData;
-  position: Position;
+  position: NodePosition;
   [k: string]: unknown;
 }
 export interface WorkspaceNodeData {

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

@@ -46,17 +46,19 @@ def register(env: str, cache: bool = True):
     ops.EXECUTORS[env] = lambda ws: execute(ws, ops.CATALOGS[env], cache=cache)
 
 
-def get_stages(ws, catalog):
+def get_stages(ws: workspace.Workspace, catalog: dict[str, ops.Op]):
     """Inputs on top/bottom are batch inputs. We decompose the graph into a DAG of components along these edges."""
     nodes = {n.id: n for n in ws.nodes}
     batch_inputs = {}
     inputs = {}
+    # For each edge in the workspacce, we record the inputs (sources)
+    # required for each node (target).
     for edge in ws.edges:
         inputs.setdefault(edge.target, []).append(edge.source)
         node = nodes[edge.target]
         op = catalog[node.data.title]
         i = op.inputs[edge.targetHandle]
-        if i.position in "top or bottom":
+        if i.side in [ops.Side.TOP, ops.Side.BOTTOM]:
             batch_inputs.setdefault(edge.target, []).append(edge.source)
     stages = []
     for bt, bss in batch_inputs.items():
@@ -93,7 +95,7 @@ async def await_if_needed(obj):
     return obj
 
 
-async def execute(ws, catalog, cache=None):
+async def execute(ws: workspace.Workspace, catalog: dict[str, ops.Op], cache=None):
     nodes = {n.id: n for n in ws.nodes}
     contexts = {n.id: Context(node=n) for n in ws.nodes}
     edges = {n.id: [] for n in ws.nodes}
@@ -108,7 +110,12 @@ async def execute(ws, catalog, cache=None):
             node.data.error = f'Operation "{node.data.title}" not found.'
             continue
         # Start tasks for nodes that have no non-batch inputs.
-        if all([i.position in "top or bottom" for i in op.inputs.values()]):
+        if all(
+            [
+                i.side in [ops.Side.TOP, ops.Side.BOTTOM]
+                for i in op.inputs.values()
+            ]
+        ):
             tasks[node.id] = [NO_INPUT]
     batch_inputs = {}
     # Run the rest until we run out of tasks.
@@ -131,7 +138,7 @@ async def execute(ws, catalog, cache=None):
                 try:
                     inputs = []
                     for i in op.inputs.values():
-                        if i.position in "top or bottom":
+                        if i.side in [ops.Side.TOP, ops.Side.BOTTOM]:
                             assert (n, i.name) in batch_inputs, f"{i.name} is missing"
                             inputs.append(batch_inputs[(n, i.name)])
                         else:
@@ -156,16 +163,16 @@ async def execute(ws, catalog, cache=None):
                 results.extend(result)
             else:  # Finished all tasks without errors.
                 if (
-                    op.type == "visualization"
-                    or op.type == "table_view"
-                    or op.type == "image"
+                    op.view_type == ops.ViewType.VISUALIZATION
+                    or op.view_type == ops.ViewType.TABLE_VIEW
+                    or op.view_type == ops.ViewType.IMAGE
                 ):
                     data.display = results[0]
                 for edge in edges[node.id]:
                     t = nodes[edge.target]
                     op = catalog[t.data.title]
                     i = op.inputs[edge.targetHandle]
-                    if i.position in "top or bottom":
+                    if i.side in [ops.Side.TOP, ops.Side.BOTTOM]:
                         batch_inputs.setdefault(
                             (edge.target, edge.targetHandle), []
                         ).extend(results)

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

@@ -61,16 +61,23 @@ class Parameter(BaseConfig):
         return Parameter(name=name, default=default, type=type)
 
 
+class Side(enum.StrEnum):
+    LEFT = "left"
+    RIGHT = "right"
+    TOP = "top"
+    BOTTOM = "bottom"
+
+
 class Input(BaseConfig):
     name: str
     type: Type
-    position: str = "left"
+    side: Side = Side.LEFT
 
 
 class Output(BaseConfig):
     name: str
     type: Type
-    position: str = "right"
+    side: Side = Side.RIGHT
 
 
 MULTI_INPUT = Input(name="multi", type="*")
@@ -84,13 +91,22 @@ def basic_outputs(*names):
     return {name: Output(name=name, type=None) for name in names}
 
 
+class ViewType(enum.StrEnum):
+    """Represents the visualization options for an operation."""
+
+    BASIC = "basic"
+    VISUALIZATION = "visualization"
+    IMAGE = "image"
+    TABLE_VIEW = "table_view"
+
+
 class Op(BaseConfig):
     func: typing.Callable = pydantic.Field(exclude=True)
     name: str
     params: dict[str, Parameter]
     inputs: dict[str, Input]
     outputs: dict[str, Output]
-    type: str = "basic"  # The UI to use for this operation.
+    view_type: ViewType = ViewType.BASIC  # The UI to use for this operation.
 
     def __call__(self, *inputs, **params):
         # Convert parameters.
@@ -133,7 +149,7 @@ def op(env: str, name: str, *, view="basic", outputs=None):
             params=params,
             inputs=inputs,
             outputs=_outputs,
-            type=view,
+            view_type=view,
         )
         CATALOGS.setdefault(env, {})
         CATALOGS[env][name] = op
@@ -143,25 +159,25 @@ def op(env: str, name: str, *, view="basic", outputs=None):
     return decorator
 
 
-def input_position(**kwargs):
-    """Decorator for specifying unusual positions for the inputs."""
+def input_side(**kwargs):
+    """Decorator for specifying unusual sides for the inputs."""
 
     def decorator(func):
         op = func.__op__
         for k, v in kwargs.items():
-            op.inputs[k].position = v
+            op.inputs[k].side = v
         return func
 
     return decorator
 
 
-def output_position(**kwargs):
-    """Decorator for specifying unusual positions for the outputs."""
+def output_side(**kwargs):
+    """Decorator for specifying unusual sides for the outputs."""
 
     def decorator(func):
         op = func.__op__
         for k, v in kwargs.items():
-            op.outputs[k].position = v
+            op.outputs[k].side = v
         return func
 
     return decorator
@@ -173,7 +189,13 @@ def no_op(*args, **kwargs):
     return None
 
 
-def register_passive_op(env: str, name: str, inputs=[], outputs=["output"], params=[]):
+def register_passive_op(
+    env: str,
+    name: str,
+    inputs: list[Input] = [],
+    outputs: list[Output] = ["output"],
+    params: list[Parameter] = [],
+):
     """A passive operation has no associated code."""
     op = Op(
         func=no_op,
@@ -209,16 +231,3 @@ def op_registration(env: str):
 
 def passive_op_registration(env: str):
     return functools.partial(register_passive_op, env)
-
-
-def register_area(env, name, params=[]):
-    """A node that represents an area. It can contain other nodes, but does not restrict movement in any way."""
-    op = Op(
-        func=no_op,
-        name=name,
-        params={p.name: p for p in params},
-        inputs={},
-        outputs={},
-        type="area",
-    )
-    CATALOGS[env][name] = op

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

@@ -14,7 +14,8 @@ class BaseConfig(pydantic.BaseModel):
     )
 
 
-class Position(BaseConfig):
+# TODO: Rename this to coordinates
+class NodePosition(BaseConfig):
     x: float
     y: float
 
@@ -29,10 +30,12 @@ class WorkspaceNodeData(BaseConfig):
 
 
 class WorkspaceNode(BaseConfig):
+    # The naming of these attributes matches the ones for the NodeBase type in React flow
+    # modyfing them will break the frontend.
     id: str
     type: str
     data: WorkspaceNodeData
-    position: Position
+    position: NodePosition
 
 
 class WorkspaceEdge(BaseConfig):
@@ -44,6 +47,13 @@ class WorkspaceEdge(BaseConfig):
 
 
 class Workspace(BaseConfig):
+    """A workspace is a representation of a computational graph that consists of nodes and edges.
+
+    Each node represents an operation or task, and the edges represent the flow of data between
+    the nodes. Each workspace is associated with an environment, which determines the operations
+    that can be performed in the workspace and the execution method for the operations.
+    """
+
     env: str = ""
     nodes: list[WorkspaceNode] = dataclasses.field(default_factory=list)
     edges: list[WorkspaceEdge] = dataclasses.field(default_factory=list)
@@ -55,6 +65,7 @@ async def execute(ws: Workspace):
 
 
 def save(ws: Workspace, path: str):
+    """Persist a workspace to a local file in JSON format."""
     j = ws.model_dump_json(indent=2)
     dirname, basename = os.path.split(path)
     # Create temp file in the same directory to make sure it's on the same filesystem.
@@ -66,7 +77,17 @@ def save(ws: Workspace, path: str):
     os.replace(temp_name, path)
 
 
-def load(path: str):
+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) as f:
         j = f.read()
     ws = Workspace.model_validate_json(j)
@@ -75,19 +96,32 @@ def load(path: str):
     return ws
 
 
-def _update_metadata(ws):
-    catalog = ops.CATALOGS.get(ws.env, {})
+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.
+
+    Returns:
+        Workspace: The updated workspace object.
+    """
+    catalog: dict[str, ops.Op] = ops.CATALOGS.get(ws.env, {})
     nodes = {node.id: node for node in ws.nodes}
     done = set()
     while len(done) < len(nodes):
         for node in ws.nodes:
             if node.id in done:
+                # TODO: Can nodes with the same ID reference different operations?
                 continue
             data = node.data
             op = catalog.get(data.title)
             if op:
                 data.meta = op
-                node.type = op.type
+                node.type = op.view_type
                 if data.error == "Unknown operation.":
                     data.error = None
             else:

lynxkite-core/tests/test_ops.py

@@ -0,0 +1,89 @@
+import inspect
+from lynxkite.core import ops
+import enum
+
+
+def test_op_decorator_no_params_no_types_default_sides():
+    @ops.op(env="test", name="add", view=ops.ViewType.BASIC, outputs=["result"])
+    def add(a, b):
+        return a + b
+
+    assert add.__op__.name == "add"
+    assert add.__op__.params == {}
+    assert add.__op__.inputs == {
+        "a": ops.Input(name="a", type=inspect._empty, side=ops.Side.LEFT),
+        "b": ops.Input(name="b", type=inspect._empty, side=ops.Side.LEFT),
+    }
+    assert add.__op__.outputs == {
+        "result": ops.Output(name="result", type=None, side=ops.Side.RIGHT)
+    }
+    assert add.__op__.view_type == ops.ViewType.BASIC
+    assert ops.CATALOGS["test"]["add"] == add.__op__
+
+
+def test_op_decorator_custom_sides():
+    @ops.input_side(a=ops.Side.RIGHT, b=ops.Side.TOP)
+    @ops.output_side(result=ops.Side.BOTTOM)
+    @ops.op(env="test", name="add", view=ops.ViewType.BASIC, outputs=["result"])
+    def add(a, b):
+        return a + b
+
+    assert add.__op__.name == "add"
+    assert add.__op__.params == {}
+    assert add.__op__.inputs == {
+        "a": ops.Input(name="a", type=inspect._empty, side=ops.Side.RIGHT),
+        "b": ops.Input(name="b", type=inspect._empty, side=ops.Side.TOP),
+    }
+    assert add.__op__.outputs == {
+        "result": ops.Output(name="result", type=None, side=ops.Side.BOTTOM)
+    }
+    assert add.__op__.view_type == ops.ViewType.BASIC
+    assert ops.CATALOGS["test"]["add"] == add.__op__
+
+
+def test_op_decorator_with_params_and_types_():
+    @ops.op(env="test", name="multiply", view=ops.ViewType.BASIC, outputs=["result"])
+    def multiply(a: int, b: float = 2.0, *, param: str = "param"):
+        return a * b
+
+    assert multiply.__op__.name == "multiply"
+    assert multiply.__op__.params == {
+        "param": ops.Parameter(name="param", default="param", type=str)
+    }
+    assert multiply.__op__.inputs == {
+        "a": ops.Input(name="a", type=int, side=ops.Side.LEFT),
+        "b": ops.Input(name="b", type=float, side=ops.Side.LEFT),
+    }
+    assert multiply.__op__.outputs == {
+        "result": ops.Output(name="result", type=None, side=ops.Side.RIGHT)
+    }
+    assert multiply.__op__.view_type == ops.ViewType.BASIC
+    assert ops.CATALOGS["test"]["multiply"] == multiply.__op__
+
+
+def test_op_decorator_with_complex_types():
+    class Color(enum.Enum):
+        RED = 1
+        GREEN = 2
+        BLUE = 3
+
+    @ops.op(env="test", name="color_op", view=ops.ViewType.BASIC, outputs=["result"])
+    def complex_op(color: Color, color_list: list[Color], color_dict: dict[str, Color]):
+        return color.name
+
+    assert complex_op.__op__.name == "color_op"
+    assert complex_op.__op__.params == {}
+    assert complex_op.__op__.inputs == {
+        "color": ops.Input(name="color", type=Color, side=ops.Side.LEFT),
+        "color_list": ops.Input(
+            name="color_list", type=list[Color], side=ops.Side.LEFT
+        ),
+        "color_dict": ops.Input(
+            name="color_dict", type=dict[str, Color], side=ops.Side.LEFT
+        ),
+    }
+    assert complex_op.__op__.view_type == ops.ViewType.BASIC
+    assert complex_op.__op__.outputs == {
+        "result": ops.Output(name="result", type=None, side=ops.Side.RIGHT)
+    }
+    assert ops.CATALOGS["test"]["color_op"] == complex_op.__op__

lynxkite-core/tests/test_workspace.py

@@ -0,0 +1,115 @@
+import os
+import pytest
+import tempfile
+from lynxkite.core import workspace
+from lynxkite.core import ops
+
+
+def test_save_load():
+    ws = workspace.Workspace(env="test")
+    ws.nodes.append(
+        workspace.WorkspaceNode(
+            id="1",
+            type="node_type",
+            data=workspace.WorkspaceNodeData(title="Node 1", params={}),
+            position=workspace.NodePosition(x=0, y=0),
+        )
+    )
+    ws.nodes.append(
+        workspace.WorkspaceNode(
+            id="2",
+            type="node_type",
+            data=workspace.WorkspaceNodeData(title="Node 2", params={}),
+            position=workspace.NodePosition(x=0, y=0),
+        )
+    )
+    ws.edges.append(
+        workspace.WorkspaceEdge(
+            id="edge1",
+            source="1",
+            target="2",
+            sourceHandle="",
+            targetHandle="",
+        )
+    )
+    path = os.path.join(tempfile.gettempdir(), "test_workspace.json")
+
+    try:
+        workspace.save(ws, path)
+        assert os.path.exists(path)
+        loaded_ws = 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)
+        sorted_ws_nodes = sorted(ws.nodes, key=lambda x: x.id)
+        sorted_loaded_ws_nodes = sorted(loaded_ws.nodes, key=lambda x: x.id)
+        # We do manual assertion on each attribute because metadata is added at
+        # loading time, which makes the objects different.
+        for node, loaded_node in zip(sorted_ws_nodes, sorted_loaded_ws_nodes):
+            assert node.id == loaded_node.id
+            assert node.type == loaded_node.type
+            assert node.data.title == loaded_node.data.title
+            assert node.data.params == loaded_node.data.params
+            assert node.position.x == loaded_node.position.x
+            assert node.position.y == loaded_node.position.y
+        sorted_ws_edges = sorted(ws.edges, key=lambda x: x.id)
+        sorted_loaded_ws_edges = sorted(loaded_ws.edges, key=lambda x: x.id)
+        for edge, loaded_edge in zip(sorted_ws_edges, sorted_loaded_ws_edges):
+            assert edge.id == loaded_edge.id
+            assert edge.source == loaded_edge.source
+            assert edge.target == loaded_edge.target
+            assert edge.sourceHandle == loaded_edge.sourceHandle
+            assert edge.targetHandle == loaded_edge.targetHandle
+    finally:
+        os.remove(path)
+
+
+@pytest.fixture(scope="session", autouse=True)
+def populate_ops_catalog():
+    ops.register_passive_op(
+        env="test",
+        name="Test Operation",
+        inputs=[],
+        params=[
+            ops.Parameter(name="param_int", default=1),
+            ops.Parameter(name="param_str", default="test"),
+        ],
+    )
+
+
+def test_update_metadata():
+    ws = workspace.Workspace(env="test")
+    ws.nodes.append(
+        workspace.WorkspaceNode(
+            id="1",
+            type="basic",
+            data=workspace.WorkspaceNodeData(
+                title="Test Operation", params={"param_int": 1}
+            ),
+            position=workspace.NodePosition(x=0, y=0),
+        )
+    )
+    ws.nodes.append(
+        workspace.WorkspaceNode(
+            id="2",
+            type="basic",
+            data=workspace.WorkspaceNodeData(title="Unknown Operation", params={}),
+            position=workspace.NodePosition(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 updated_ws.nodes[0].data.params == {"param_int": 1}
+    assert updated_ws.nodes[0].data.meta.params == {
+        "param_int": ops.Parameter(name="param_int", default=1),
+        "param_str": ops.Parameter(name="param_str", default="test"),
+    }
+    assert not hasattr(updated_ws.nodes[1].data, "meta")
+    assert updated_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

lynxkite-graph-analytics/src/lynxkite_plugins/graph_analytics/lynxkite_ops.py

@@ -1,7 +1,7 @@
 """Graph analytics operations. To be split into separate files when we have more."""
 
 import os
-from lynxkite.core import ops
+from lynxkite.core import ops, workspace
 from collections import deque
 import dataclasses
 import functools
@@ -34,6 +34,7 @@ class RelationDefinition:
     target_key: str  # The column in the target table that contains the node ID.
 
 
+# TODO: Convert this to Pydantic
 @dataclasses.dataclass
 class Bundle:
     """A collection of DataFrames and other data.
@@ -115,32 +116,50 @@ def disambiguate_edges(ws):
 
 
 @ops.register_executor(ENV)
-async def execute(ws):
-    catalog = ops.CATALOGS[ENV]
+async def execute(ws: workspace.Workspace):
+    catalog: dict[str, ops.Op] = ops.CATALOGS[ENV]
     disambiguate_edges(ws)
-    outputs = {}
+    computed_outputs = {}
     failed = 0
-    while len(outputs) + failed < len(ws.nodes):
+    while len(computed_outputs) + failed < len(ws.nodes):
         for node in ws.nodes:
-            if node.id in outputs:
+            if node.id in computed_outputs:
                 continue
             # TODO: Take the input/output handles into account.
-            inputs = [edge.source for edge in ws.edges if edge.target == node.id]
-            if all(input in outputs for input in inputs):
-                inputs = [outputs[input] for input in inputs]
+            operation_inputs = [
+                edge.source for edge in ws.edges if edge.target == node.id
+            ]
+            if all(input in computed_outputs for input in operation_inputs):
+                # All inputs for this node are ready, we can compute the output.
+                operation_inputs = [
+                    computed_outputs[input] for input in operation_inputs
+                ]
                 data = node.data
-                op = catalog[data.title]
                 params = {**data.params}
-                # Convert inputs.
                 try:
-                    for i, (x, p) in enumerate(zip(inputs, op.inputs.values())):
-                        if p.type == nx.Graph and isinstance(x, Bundle):
-                            inputs[i] = x.to_nx()
-                        elif p.type == Bundle and isinstance(x, nx.Graph):
-                            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)
+                    op = catalog[data.title]
+                    # Convert inputs types  to match operation signature.
+                    for i, (input_value, input_signature) in enumerate(
+                        zip(operation_inputs, op.inputs.values())
+                    ):
+                        if input_signature.type == nx.Graph and isinstance(
+                            input_value, Bundle
+                        ):
+                            operation_inputs[i] = input_value.to_nx()
+                        elif input_signature.type == Bundle and isinstance(
+                            input_value, nx.Graph
+                        ):
+                            operation_inputs[i] = Bundle.from_nx(input_value)
+                        elif input_signature.type == Bundle and isinstance(
+                            input_value, pd.DataFrame
+                        ):
+                            operation_inputs[i] = Bundle.from_df(input_value)
+                    output = op(*operation_inputs, **params)
+                except KeyError:
+                    traceback.print_exc()
+                    data.error = "Operation not found in catalog"
+                    failed += 1
+                    continue
                 except Exception as e:
                     traceback.print_exc()
                     data.error = str(e)
@@ -148,13 +167,16 @@ async def execute(ws):
                     continue
                 if len(op.inputs) == 1 and op.inputs.get("multi") == "*":
                     # It's a flexible input. Create n+1 handles.
-                    data.inputs = {f"input{i}": None for i in range(len(inputs) + 1)}
+                    # TODO: How is this used? Why we define the inputs in the WorkspaceNodeData?
+                    data.inputs = {
+                        f"input{i}": None for i in range(len(operation_inputs) + 1)
+                    }
                 data.error = None
-                outputs[node.id] = output
+                computed_outputs[node.id] = output
                 if (
-                    op.type == "visualization"
-                    or op.type == "table_view"
-                    or op.type == "image"
+                    op.view_type == ops.ViewType.VISUALIZATION
+                    or op.view_type == ops.ViewType.TABLE_VIEW
+                    or op.view_type == ops.ViewType.IMAGE
                 ):
                     data.display = output
 
@@ -188,6 +210,7 @@ def create_scale_free_graph(*, nodes: int = 10):
 @op("Compute PageRank")
 @nx_node_attribute_func("pagerank")
 def compute_pagerank(graph: nx.Graph, *, damping=0.85, iterations=100):
+    # TODO: This requires scipy to be installed.
     return nx.pagerank(graph, alpha=damping, max_iter=iterations)
 
 
@@ -281,7 +304,7 @@ def _map_color(value):
         ]
 
 
-@op("Visualize graph", view="visualization")
+@op("Visualize graph", view=ops.ViewType.VISUALIZATION)
 def visualize_graph(graph: Bundle, *, color_nodes_by: ops.NodeAttribute = None):
     nodes = graph.dfs["nodes"].copy()
     if color_nodes_by:
@@ -335,7 +358,7 @@ def collect(df: pd.DataFrame):
     return df.values.tolist()
 
 
-@op("View tables", view="table_view")
+@op("View tables", view=ops.ViewType.TABLE_VIEW)
 def view_tables(bundle: Bundle):
     v = {
         "dataframes": {

lynxkite-graph-analytics/src/lynxkite_plugins/graph_analytics/networkx_ops.py

@@ -54,7 +54,7 @@ def register_networkx(env: str):
                 params=params,
                 inputs=inputs,
                 outputs={"output": ops.Output(name="output", type=nx.Graph)},
-                type="basic",
+                view_type=ops.ViewType.BASIC,
             )
             cat[name] = op
 

lynxkite-graph-analytics/src/lynxkite_plugins/graph_analytics/pytorch_model_ops.py

@@ -13,10 +13,10 @@ def reg(name, inputs=[], outputs=None, params=[]):
         ENV,
         name,
         inputs=[
-            ops.Input(name=name, position="bottom", type="tensor") for name in inputs
+            ops.Input(name=name, side=ops.Side.BOTTOM, type="tensor") for name in inputs
         ],
         outputs=[
-            ops.Output(name=name, position="top", type="tensor") for name in outputs
+            ops.Output(name=name, side=ops.Side.TOP, type="tensor") for name in outputs
         ],
         params=params,
     )
@@ -64,6 +64,4 @@ reg(
         ),
         P.basic("lr", 0.001),
     ],
-)
-
-ops.register_area(ENV, "Repeat", params=[ops.Parameter.basic("times", 1, int)])
+)

lynxkite-graph-analytics/tests/test_lynxkite_ops.py

@@ -0,0 +1,100 @@
+import sys
+
+# Add the project root to sys.path
+sys.path.insert(0, "/home/chema/work/lynxkite-2024/lynxkite-graph-analytics")
+import pandas as pd
+import pytest
+import networkx as nx
+
+from lynxkite.core import workspace
+
+# from lynxkite_plugins.graph_analytics.lynxkite_ops import execute
+from src.lynxkite_plugins.graph_analytics.lynxkite_ops import Bundle, execute, op
+
+
+async def test_execute_operation_not_in_catalog():
+    ws = workspace.Workspace(env="test")
+    ws.nodes.append(
+        workspace.WorkspaceNode(
+            id="1",
+            type="node_type",
+            data=workspace.WorkspaceNodeData(title="Non existing op", params={}),
+            position=workspace.NodePosition(x=0, y=0),
+        )
+    )
+    await execute(ws)
+    assert ws.nodes[0].data.error == "Operation not found in catalog"
+
+
+async def test_execute_operation_inputs_correct_cast():
+    # Test that the automatic casting of operation inputs works correctly.
+
+    @op("Create Bundle")
+    def create_bundle() -> Bundle:
+        df = pd.DataFrame({"source": [1, 2, 3], "target": [4, 5, 6]})
+        return Bundle(dfs={"edges": df})
+
+    @op("Bundle to Graph")
+    def bundle_to_graph(graph: nx.Graph) -> nx.Graph:
+        return graph
+
+    @op("Graph to Bundle")
+    def graph_to_bundle(bundle: Bundle) -> pd.DataFrame:
+        return list(bundle.dfs.values())[0]
+
+    @op("Dataframe to Bundle")
+    def dataframe_to_bundle(bundle: Bundle) -> Bundle:
+        return bundle
+
+    ws = workspace.Workspace(env="test")
+    ws.nodes.append(
+        workspace.WorkspaceNode(
+            id="1",
+            type="node_type",
+            data=workspace.WorkspaceNodeData(title="Create Bundle", params={}),
+            position=workspace.NodePosition(x=0, y=0),
+        )
+    )
+    ws.nodes.append(
+        workspace.WorkspaceNode(
+            id="2",
+            type="node_type",
+            data=workspace.WorkspaceNodeData(title="Bundle to Graph", params={}),
+            position=workspace.NodePosition(x=100, y=0),
+        )
+    )
+    ws.nodes.append(
+        workspace.WorkspaceNode(
+            id="3",
+            type="node_type",
+            data=workspace.WorkspaceNodeData(title="Graph to Bundle", params={}),
+            position=workspace.NodePosition(x=200, y=0),
+        )
+    )
+    ws.nodes.append(
+        workspace.WorkspaceNode(
+            id="4",
+            type="node_type",
+            data=workspace.WorkspaceNodeData(title="Dataframe to Bundle", params={}),
+            position=workspace.NodePosition(x=300, y=0),
+        )
+    )
+    ws.edges = [
+        workspace.WorkspaceEdge(
+            id="1", source="1", target="2", sourceHandle="1", targetHandle="2"
+        ),
+        workspace.WorkspaceEdge(
+            id="2", source="2", target="3", sourceHandle="2", targetHandle="3"
+        ),
+        workspace.WorkspaceEdge(
+            id="3", source="3", target="4", sourceHandle="3", targetHandle="4"
+        ),
+    ]
+
+    await execute(ws)
+
+    assert all([node.data.error is None for node in ws.nodes])
+
+
+if __name__ == "__main__":
+    pytest.main()

lynxkite-lynxscribe/src/lynxkite_plugins/lynxscribe/llm_ops.py

@@ -68,13 +68,13 @@ def split_document(input, *, delimiter: str = "\\n\\n"):
     return pd.DataFrame(chunks, columns=["text"])
 
 
-@ops.input_position(input="top")
+@ops.input_side(input=ops.Side.TOP)
 @op("Build document graph")
 def build_document_graph(input):
     return [{"source": i, "target": i + 1} for i in range(len(input) - 1)]
 
 
-@ops.input_position(nodes="top", edges="top")
+@ops.input_side(nodes=ops.Side.TOP, edges=ops.Side.TOP)
 @op("Predict links")
 def predict_links(nodes, edges):
     """A placeholder for a real algorithm. For now just adds 2-hop neighbors."""
@@ -89,7 +89,7 @@ def predict_links(nodes, edges):
     return edges + new_edges
 
 
-@ops.input_position(nodes="top", edges="top")
+@ops.input_side(nodes=ops.Side.TOP, edges=ops.Side.TOP)
 @op("Add neighbors")
 def add_neighbors(nodes, edges, item):
     nodes = pd.DataFrame(nodes)
@@ -133,7 +133,7 @@ def ask_llm(input, *, model: str, accepted_regex: str = None, max_tokens: int =
     return [{**input, "response": r} for r in results]
 
 
-@op("View", view="table_view")
+@op("View", view=ops.ViewType.TABLE_VIEW)
 def view(input, *, _ctx: one_by_one.Context):
     v = _ctx.last_result
     if v:
@@ -152,8 +152,8 @@ def view(input, *, _ctx: one_by_one.Context):
     return v
 
 
-@ops.input_position(input="right")
-@ops.output_position(output="left")
+@ops.input_side(input=ops.Side.RIGHT)
+@ops.output_side(output=ops.Side.LEFT)
 @op("Loop")
 def loop(input, *, max_iterations: int = 3, _ctx: one_by_one.Context):
     """Data can flow back here max_iterations-1 times."""
@@ -174,7 +174,7 @@ class RagEngine(enum.Enum):
     Custom = "Custom"
 
 
-@ops.input_position(db="top")
+@ops.input_side(db=ops.Side.TOP)
 @op("RAG")
 def rag(
     input,

lynxkite-lynxscribe/src/lynxkite_plugins/lynxscribe/lynxscribe_ops.py

@@ -24,7 +24,7 @@ from lynxkite.core.executors import one_by_one
 ENV = "LynxScribe"
 one_by_one.register(ENV)
 op = ops.op_registration(ENV)
-output_on_top = ops.output_position(output="top")
+output_on_top = ops.output_side(output=ops.Side.TOP)
 
 
 @output_on_top
@@ -42,7 +42,7 @@ def llm(*, name="openai"):
 
 
 @output_on_top
-@ops.input_position(llm="bottom")
+@ops.input_side(llm=ops.Side.BOTTOM)
 @op("Text embedder")
 def text_embedder(llm, *, model="text-embedding-ada-002"):
     llm = llm[0]["llm"]
@@ -51,7 +51,7 @@ def text_embedder(llm, *, model="text-embedding-ada-002"):
 
 
 @output_on_top
-@ops.input_position(vector_store="bottom", text_embedder="bottom")
+@ops.input_side(vector_store=ops.Side.BOTTOM, text_embedder=ops.Side.BOTTOM)
 @op("RAG graph")
 def rag_graph(vector_store, text_embedder):
     vector_store = vector_store[0]["vector_store"]
@@ -78,7 +78,7 @@ DEFAULT_NEGATIVE_ANSWER = "I'm sorry, but the data I've been trained on does not
 
 
 @output_on_top
-@ops.input_position(rag_graph="bottom", scenario_selector="bottom", llm="bottom")
+@ops.input_side(rag_graph=ops.Side.BOTTOM, scenario_selector=ops.Side.BOTTOM, llm=ops.Side.BOTTOM)
 @op("RAG chatbot")
 def rag_chatbot(
     rag_graph,
@@ -107,7 +107,7 @@ def rag_chatbot(
 
 
 @output_on_top
-@ops.input_position(processor="bottom")
+@ops.input_side(processor=ops.Side.BOTTOM)
 @op("Chat processor")
 def chat_processor(processor, *, _ctx: one_by_one.Context):
     cfg = _ctx.last_result or {
@@ -152,7 +152,7 @@ def mask(*, name="", regex="", exceptions="", mask_pattern=""):
     }
 
 
-@ops.input_position(chat_api="bottom")
+@ops.input_side(chat_api=ops.Side.BOTTOM)
 @op("Test Chat API")
 async def test_chat_api(message, chat_api, *, show_details=False):
     chat_api = chat_api[0]["chat_api"]
@@ -173,7 +173,7 @@ def input_chat(*, chat: str):
 
 
 @output_on_top
-@ops.input_position(chatbot="bottom", chat_processor="bottom", knowledge_base="bottom")
+@ops.input_side(chatbot=ops.Side.BOTTOM, chat_processor=ops.Side.BOTTOM, knowledge_base=ops.Side.BOTTOM)
 @op("Chat API")
 def chat_api(chatbot, chat_processor, knowledge_base, *, model="gpt-4o-mini"):
     chatbot = chatbot[0]["chatbot"]
@@ -205,7 +205,7 @@ def knowledge_base(
     }
 
 
-@op("View", view="table_view")
+@op("View", view=ops.ViewType.TABLE_VIEW)
 def view(input):
     columns = [str(c) for c in input.keys() if not str(c).startswith("_")]
     v = {

lynxkite-lynxscribe/tests/test_llm_ops.py

@@ -8,7 +8,7 @@ def make_node(id, op, type="basic", **params):
     return workspace.WorkspaceNode(
         id=id,
         type=type,
-        position=workspace.Position(x=0, y=0),
+        position=workspace.NodePosition(x=0, y=0),
         data=workspace.WorkspaceNodeData(title=op, params=params),
     )
 
@@ -43,7 +43,7 @@ class LLMOpsTest(unittest.IsolatedAsyncioTestCase):
                     filename="/Users/danieldarabos/Downloads/aimo-train.csv",
                     key="problem",
                 ),
-                make_node("1", "View", type="table_view"),
+                make_node("1", "View", type=ops.ViewType.TABLE_VIEW),
             ],
             edges=[make_edge("0", "1")],
         )

lynxkite-pillow-example/src/lynxkite_plugins/pillow_example/__init__.py

@@ -56,7 +56,7 @@ def to_grayscale(image: Image):
     return image.convert("L")
 
 
-@op("View image", view="image")
+@op("View image", view=ops.ViewType.IMAGE)
 def view_image(image: Image):
     buffered = io.BytesIO()
     image.save(buffered, format="JPEG")