Use Pydantic for op metadata.

server/{basic_ops.py → lynxkite_ops.py}

@@ -20,6 +20,12 @@ def compute_pagerank(graph: nx.Graph, *, damping=0.85, iterations=100):
   return nx.pagerank(graph, alpha=damping, max_iter=iterations)
 
 
+@ops.op("Sample graph")
+def create_scale_free_graph(*, nodes: int = 10):
+  '''Creates a scale-free graph with the given number of nodes.'''
+  return nx.scale_free_graph(nodes)
+
+
 def _map_color(value):
   cmap = matplotlib.cm.get_cmap('viridis')
   value = (value - value.min()) / (value.max() - value.min())

server/main.py

@@ -3,17 +3,17 @@ import fastapi
 import pathlib
 from . import ops
 from . import workspace
-from . import basic_ops
+from . import lynxkite_ops
 # from . import networkx_ops
-# from . import pytorch_model_ops
-from . import lynxscribe_ops
+from . import pytorch_model_ops
+# from . import lynxscribe_ops
 
 app = fastapi.FastAPI()
 
 
 @app.get("/api/catalog")
 def get_catalog():
-    return [op.to_json() for op in ops.ALL_OPS.values()]
+    return [op.model_dump() for op in ops.ALL_OPS.values()]
 
 
 class SaveRequest(workspace.BaseConfig):

server/networkx_ops.py

@@ -26,7 +26,7 @@ for (name, func) in nx.__dict__.items():
     sig = inspect.signature(func)
     inputs = {k: nx.Graph for k in func.graphs}
     params = {
-      name: ops.Parameter(
+      name: ops.Parameter.basic(
           name, str(param.default)
         if type(param.default) in [str, int, float]
         else None,

server/ops.py

@@ -1,60 +1,64 @@
 '''API for implementing LynxKite operations.'''
+from __future__ import annotations
 import dataclasses
 import enum
 import functools
 import inspect
 import networkx as nx
 import pandas as pd
+import pydantic
 import typing
+from typing_extensions import Annotated
 
 ALL_OPS = {}
-PARAM_TYPE = type[typing.Any]
 typeof = type # We have some arguments called "type".
-
-@dataclasses.dataclass
-class Parameter:
+def type_to_json(t):
+  if isinstance(t, type) and issubclass(t, enum.Enum):
+    return {'enum': list(t.__members__.keys())}
+  if isinstance(t, tuple) and t[0] == 'collapsed':
+    return {'collapsed': str(t[1])}
+  return {'type': str(t)}
+Type = Annotated[
+  typing.Any, pydantic.PlainSerializer(type_to_json, return_type=dict)
+]
+class BaseConfig(pydantic.BaseModel):
+  model_config = pydantic.ConfigDict(
+    arbitrary_types_allowed=True,
+  )
+
+
+class Parameter(BaseConfig):
   '''Defines a parameter for an operation.'''
   name: str
   default: any
-  type: PARAM_TYPE = None
+  type: Type = None
 
   @staticmethod
   def options(name, options, default=None):
     e = enum.Enum(f'OptionsFor_{name}', options)
-    return Parameter(name, e[default or options[0]], e)
+    return Parameter.basic(name, e[default or options[0]], e)
 
   @staticmethod
   def collapsed(name, default, type=None):
-    return Parameter(name, default, ('collapsed', type or typeof(default)))
-
-  def __post_init__(self):
-    if self.default is inspect._empty:
-      self.default = None
-    if self.type is None or self.type is inspect._empty:
-      self.type = type(self.default)
-  def to_json(self):
-    t = str(self.type)
-    default = self.default
-    if isinstance(self.type, type) and issubclass(self.type, enum.Enum):
-      t = {'enum': list(self.type.__members__.keys())}
-      default = self.default.name if self.default else t['enum'][0]
-    if isinstance(self.type, tuple) and self.type[0] == 'collapsed':
-      t = {'collapsed': str(self.type[1])}
-    return {
-      'name': self.name,
-      'default': default,
-      'type': t,
-    }
+    return Parameter.basic(name, default, ('collapsed', type or typeof(default)))
 
-@dataclasses.dataclass
-class Op:
-  func: callable
+  @staticmethod
+  def basic(name, default=None, type=None):
+    if default is inspect._empty:
+      default = None
+    if type is None or type is inspect._empty:
+      type = typeof(default) if default else None
+    return Parameter(name=name, default=default, type=type)
+
+
+class Op(BaseConfig):
+  func: callable = pydantic.Field(exclude=True)
   name: str
   params: dict[str, Parameter]
-  inputs: dict # name -> type
-  outputs: dict # name -> type
+  inputs: dict[str, Type] # name -> type
+  outputs: dict[str, Type] # name -> type
   type: str # The UI to use for this operation.
-  sub_nodes: list = None # If set, these nodes can be placed inside the operation's node.
+  sub_nodes: list[Op] = None # If set, these nodes can be placed inside the operation's node.
 
   def __call__(self, *inputs, **params):
     # Convert parameters.
@@ -74,18 +78,6 @@ class Op:
     res = self.func(*inputs, **params)
     return res
 
-  def to_json(self):
-    return {
-      'type': self.type,
-      'data': {
-        'title': self.name,
-        'inputs': {i: str(type) for i, type in self.inputs.items()},
-        'outputs': {o: str(type) for o, type in self.outputs.items()},
-        'params': [p.to_json() for p in self.params.values()],
-      },
-      'sub_nodes': [sub.to_json() for sub in self.sub_nodes.values()] if self.sub_nodes else None,
-    }
-
 
 @dataclasses.dataclass
 class RelationDefinition:
@@ -105,9 +97,9 @@ class Bundle:
   Can efficiently represent a knowledge graph (homogeneous or heterogeneous) or tabular data.
   It can also carry other data, such as a trained model.
   '''
-  dfs: dict = dataclasses.field(default_factory=dict) # name -> DataFrame
+  dfs: dict[str, pd.DataFrame] = dataclasses.field(default_factory=dict)
   relations: list[RelationDefinition] = dataclasses.field(default_factory=list)
-  other: dict = None
+  other: dict[str, typing.Any] = None
 
   @classmethod
   def from_nx(cls, graph: nx.Graph):
@@ -161,9 +153,9 @@ def op(name, *, view='basic', sub_nodes=None):
     params = {}
     for n, param in sig.parameters.items():
       if param.kind == param.KEYWORD_ONLY:
-        params[n] = Parameter(n, param.default, param.annotation)
+        params[n] = Parameter.basic(n, param.default, param.annotation)
     outputs = {'output': 'yes'} if view == 'basic' else {} # Maybe more fancy later.
-    op = Op(func, name, params=params, inputs=inputs, outputs=outputs, type=view)
+    op = Op(func=func, name=name, params=params, inputs=inputs, outputs=outputs, type=view)
     if sub_nodes is not None:
       op.sub_nodes = sub_nodes
       op.type = 'sub_flow'

server/pytorch_model_ops.py

@@ -23,11 +23,11 @@ def register_layer(name):
       for name, param in sig.parameters.items()
       if param.kind != param.KEYWORD_ONLY}
     params = {
-      name: ops.Parameter(name, param.default, param.annotation)
+      name: ops.Parameter.basic(name, param.default, param.annotation)
       for name, param in sig.parameters.items()
       if param.kind == param.KEYWORD_ONLY}
     outputs = {'x': 'tensor'}
-    LAYERS[name] = ops.Op(func, name, params=params, inputs=inputs, outputs=outputs, type='vertical')
+    LAYERS[name] = ops.Op(func=func, name=name, params=params, inputs=inputs, outputs=outputs, type='vertical')
     return func
   return decorator
 
@@ -64,7 +64,7 @@ def nonlinearity(x, *, type: Nonlinearity):
 
 def register_area(name, params=[]):
   '''A node that represents an area. It can contain other nodes, but does not restrict movement in any way.'''
-  op = ops.Op(ops.no_op, name, params={p.name: p for p in params}, inputs={}, outputs={}, type='area')
+  op = ops.Op(func=ops.no_op, name=name, params={p.name: p for p in params}, inputs={}, outputs={}, type='area')
   LAYERS[name] = op
 
-register_area('Repeat', params=[ops.Parameter('times', 1, int)])
+register_area('Repeat', params=[ops.Parameter.basic('times', 1, int)])

web/src/LynxKiteFlow.svelte

@@ -31,7 +31,7 @@
   const backendWorkspace = useQuery(['workspace', path], async () => {
     const res = await fetch(`/api/load?path=${path}`);
     return res.json();
-  }, {staleTime: 10000});
+  }, {staleTime: 10000, retry: false});
   const mutation = useMutation(async(update) => {
     const res = await fetch('/api/save', {
       method: 'POST',
@@ -79,13 +79,20 @@
     };
   }
   function addNode(e) {
-    const node = {...e.detail};
+    const meta = {...e.detail};
     nodes.update((n) => {
+      const node = {
+        type: meta.type,
+        data: {
+          title: meta.name,
+          params: Object.fromEntries(
+            Object.values(meta.params).map((p) => [p.name, p.default])),
+          inputs: meta.inputs,
+          outputs: meta.outputs,
+        },
+      };
       node.position = screenToFlowPosition({x: nodeSearchSettings.pos.x, y: nodeSearchSettings.pos.y});
-      node.data = { ...node.data };
       const title = node.data.title;
-      node.data.params = Object.fromEntries(
-        node.data.params.map((p) => [p.name, p.default]));
       let i = 1;
       node.id = `${title} ${i}`;
       while (n.find((x) => x.id === node.id)) {

web/src/NodeSearch.svelte

@@ -9,7 +9,7 @@
   let selectedIndex = 0;
   onMount(() => searchBox.focus());
   $: fuse = new Fuse(boxes, {
-    keys: ['data.title']
+    keys: ['name']
   })
   function onInput() {
     hits = fuse.search(searchBox.value);
@@ -58,7 +58,7 @@
         on:click={addSelected}
         class="search-result"
         class:selected={index == selectedIndex}>
-        {box.item.data.title}
+        {box.item.name}
       </div>
     {/each}
   </div>