Choose environment per workspace.

server/llm_ops.py

@@ -7,6 +7,8 @@ import pandas as pd
 client = openai.OpenAI(base_url="http://localhost:11434/v1")
 CACHE = {}
 
+op = ops.op_registration('LLM logic')
+
 def chat(*args, **kwargs):
   key = json.dumps({'args': args, 'kwargs': kwargs})
   if key not in CACHE:
@@ -14,11 +16,11 @@ def chat(*args, **kwargs):
     CACHE[key] = [c.message.content for c in completion.choices]
   return CACHE[key]
 
-@ops.op("Input")
+@op("Input")
 def input(*, filename: ops.PathStr, key: str):
   return pd.read_csv(filename).rename(columns={key: 'text'})
 
-@ops.op("Create prompt")
+@op("Create prompt")
 def create_prompt(input, *, template: ops.LongStr):
   assert template, 'Please specify the template. Refer to columns using their names in uppercase.'
   df = input.copy()
@@ -32,15 +34,15 @@ def create_prompt(input, *, template: ops.LongStr):
   return df
 
 
-@ops.op("Ask LLM")
-def ask_llm(input, *, model: str, choices: list = None, max_tokens: int = 100):
+@op("Ask LLM")
+def ask_llm(input, *, model: str, accepted_regex: str = None, max_tokens: int = 100):
   assert model, 'Please specify the model.'
   assert 'prompt' in input.columns, 'Please create the prompt first.'
   df = input.copy()
   g = {}
-  if choices:
+  if accepted_regex:
     g['extra_body'] = {
-      "guided_choice": choices.split()
+      "guided_regex": accepted_regex,
     }
   for i, row in df.iterrows():
     [res] = chat(
@@ -54,7 +56,7 @@ def ask_llm(input, *, model: str, choices: list = None, max_tokens: int = 100):
     df.loc[i, 'response'] = res
   return df
 
-@ops.op("View", view="table_view")
+@op("View", view="table_view")
 def view(input):
   v = {
     'dataframes': { 'df': {
@@ -66,7 +68,7 @@ def view(input):
 
 @ops.input_position(input="right")
 @ops.output_position(output="left")
-@ops.op("Loop")
+@op("Loop")
 def loop(input, *, max_iterations: int = 10):
   '''Data can flow back here until it becomes empty or reaches the limit.'''
   return input

server/lynxkite_ops.py

@@ -3,26 +3,29 @@ from . import ops
 import matplotlib
 import networkx as nx
 import pandas as pd
+import traceback
 
-@ops.op("Import Parquet")
+op = ops.op_registration('LynxKite')
+
+@op("Import Parquet")
 def import_parquet(*, filename: str):
   '''Imports a parquet file.'''
   return pd.read_parquet(filename)
 
-@ops.op("Create scale-free graph")
+@op("Create scale-free 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)
 
-@ops.op("Compute PageRank")
+@op("Compute PageRank")
 @ops.nx_node_attribute_func('pagerank')
 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.'''
+@op("Sample graph")
+def sample_graph(graph: nx.Graph, *, nodes: int = 100):
+  '''Takes a subgraph.'''
   return nx.scale_free_graph(nodes)
 
 
@@ -32,10 +35,9 @@ def _map_color(value):
   rgba = cmap(value)
   return ['#{:02x}{:02x}{:02x}'.format(int(r*255), int(g*255), int(b*255)) for r, g, b in rgba[:, :3]]
 
-@ops.op("Visualize graph", view="visualization")
+@op("Visualize graph", view="visualization")
 def visualize_graph(graph: ops.Bundle, *, color_nodes_by: 'node_attribute' = None):
   nodes = graph.dfs['nodes'].copy()
-  node_attributes = sorted(nodes.columns)
   if color_nodes_by:
     nodes['color'] = _map_color(nodes[color_nodes_by])
   nodes = nodes.to_records()
@@ -76,7 +78,7 @@ def visualize_graph(graph: ops.Bundle, *, color_nodes_by: 'node_attribute' = Non
   }
   return v
 
-@ops.op("View tables", view="table_view")
+@op("View tables", view="table_view")
 def view_tables(bundle: ops.Bundle):
   v = {
     'dataframes': { name: {
@@ -87,3 +89,44 @@ def view_tables(bundle: ops.Bundle):
     'other': bundle.other,
   }
   return v
+
+@ops.register_executor('LynxKite')
+def execute(ws):
+    catalog = ops.CATALOGS['LynxKite']
+    # Nodes are responsible for interpreting/executing their child nodes.
+    nodes = [n for n in ws.nodes if not n.parentId]
+    children = {}
+    for n in ws.nodes:
+        if n.parentId:
+            children.setdefault(n.parentId, []).append(n)
+    outputs = {}
+    failed = 0
+    while len(outputs) + failed < len(nodes):
+        for node in nodes:
+            if node.id in outputs:
+                continue
+            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]
+                data = node.data
+                op = catalog[data.title]
+                params = {**data.params}
+                if op.sub_nodes:
+                    sub_nodes = children.get(node.id, [])
+                    sub_node_ids = [node.id for node in sub_nodes]
+                    sub_edges = [edge for edge in ws.edges if edge.source in sub_node_ids]
+                    params['sub_flow'] = {'nodes': sub_nodes, 'edges': sub_edges}
+                try:
+                  output = op(*inputs, **params)
+                except Exception as e:
+                  traceback.print_exc()
+                  data.error = str(e)
+                  failed += 1
+                  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)}
+                data.error = None
+                outputs[node.id] = output
+                if op.type == 'visualization' or op.type == 'table_view':
+                    data.view = output

server/lynxscribe_ops.py

@@ -1,6 +1,7 @@
 '''An example of passive ops. Just using LynxKite to describe the configuration of a complex system.'''
-from .ops import register_passive_op as reg, Parameter as P, MULTI_INPUT
+from .ops import passive_op_registration, Parameter as P, MULTI_INPUT
 
+reg = passive_op_registration('LynxScribe')
 reg('Scrape documents', params=[P.basic('url', '')])
 reg('Conversation logs')
 reg('Extract graph', inputs=['input'])

server/main.py

@@ -8,7 +8,7 @@ from . import workspace
 from . import lynxkite_ops
 # from . import networkx_ops
 # from . import pytorch_model_ops
-# from . import lynxscribe_ops
+from . import lynxscribe_ops
 from . import llm_ops
 
 app = fastapi.FastAPI()
@@ -16,7 +16,9 @@ app = fastapi.FastAPI()
 
 @app.get("/api/catalog")
 def get_catalog():
-    return {op.name: op.model_dump() for op in ops.ALL_OPS.values()}
+    return {
+        k: {op.name: op.model_dump() for op in v.values()}
+        for k, v in ops.CATALOGS.items()}
 
 
 class SaveRequest(workspace.BaseConfig):

server/networkx_ops.py

@@ -5,7 +5,7 @@ import inspect
 import networkx as nx
 
 
-def wrapped(func):
+def wrapped(name: str, func):
   @functools.wraps(func)
   def wrapper(*args, **kwargs):
     for k, v in kwargs.items():
@@ -20,24 +20,27 @@ def wrapped(func):
     return graph
   return wrapper
 
+def register_networkx(env: str):
+  ops.CATALOGS.setdefault(env, {})
+  for (name, func) in nx.__dict__.items():
+    if hasattr(func, 'graphs'):
+      sig = inspect.signature(func)
+      inputs = {k: nx.Graph for k in func.graphs}
+      params = {
+        name: ops.Parameter.basic(
+            name, str(param.default)
+          if type(param.default) in [str, int, float]
+          else None,
+          param.annotation)
+        for name, param in sig.parameters.items()
+        if name not in ['G', 'backend', 'backend_kwargs']}
+      for p in params.values():
+        if not p.type:
+          # Guess the type based on the name.
+          if len(p.name) == 1:
+            p.type = int
+      name = "NX › " + name.replace('_', ' ').title()
+      op = ops.Op(wrapped(name, func), name, params=params, inputs=inputs, outputs={'output': 'yes'}, type='basic')
+      ops.CATALOGS[env][name] = op
 
-for (name, func) in nx.__dict__.items():
-  if hasattr(func, 'graphs'):
-    sig = inspect.signature(func)
-    inputs = {k: nx.Graph for k in func.graphs}
-    params = {
-      name: ops.Parameter.basic(
-          name, str(param.default)
-        if type(param.default) in [str, int, float]
-        else None,
-        param.annotation)
-      for name, param in sig.parameters.items()
-      if name not in ['G', 'backend', 'backend_kwargs']}
-    for p in params.values():
-      if not p.type:
-        # Guess the type based on the name.
-        if len(p.name) == 1:
-          p.type = int
-    name = "NX › " + name.replace('_', ' ').title()
-    op = ops.Op(wrapped(func), name, params=params, inputs=inputs, outputs={'output': 'yes'}, type='basic')
-    ops.ALL_OPS[name] = op
+register_networkx('LynxKite')

server/ops.py

@@ -10,7 +10,9 @@ import pydantic
 import typing
 from typing_extensions import Annotated
 
-ALL_OPS = {}
+CATALOGS = {}
+EXECUTORS = {}
+
 typeof = type # We have some arguments called "type".
 def type_to_json(t):
   if isinstance(t, type) and issubclass(t, enum.Enum):
@@ -166,7 +168,7 @@ def nx_node_attribute_func(name):
   return decorator
 
 
-def op(name, *, view='basic', sub_nodes=None):
+def op(env: str, name: str, *, view='basic', sub_nodes=None):
   '''Decorator for defining an operation.'''
   def decorator(func):
     sig = inspect.signature(func)
@@ -184,7 +186,8 @@ def op(name, *, view='basic', sub_nodes=None):
     if sub_nodes is not None:
       op.sub_nodes = sub_nodes
       op.type = 'sub_flow'
-    ALL_OPS[name] = op
+    CATALOGS.setdefault(env, {})
+    CATALOGS[env][name] = op
     func.__op__ = op
     return func
   return decorator
@@ -212,7 +215,7 @@ def no_op(*args, **kwargs):
     return args[0]
   return Bundle()
 
-def register_passive_op(name, inputs=[], outputs=['output'], params=[]):
+def register_passive_op(env: str, name: str, inputs=[], outputs=['output'], params=[]):
   '''A passive operation has no associated code.'''
   op = Op(
     func=no_op,
@@ -224,10 +227,19 @@ def register_passive_op(name, inputs=[], outputs=['output'], params=[]):
     outputs=dict(
       (o, Output(name=o, type=None)) if isinstance(o, str)
       else (o.name, o) for o in outputs))
-  ALL_OPS[name] = op
+  CATALOGS.setdefault(env, {})
+  CATALOGS[env][name] = op
   return op
 
-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 = register_passive_op(name, params=params)
-  op.type = 'area'
+def register_executor(env: str):
+  '''Decorator for registering an executor.'''
+  def decorator(func):
+    EXECUTORS[env] = func
+    return func
+  return decorator
+
+def op_registration(env: str):
+  return functools.partial(op, env)
+
+def passive_op_registration(env: str):
+  return functools.partial(register_passive_op, env)

server/pytorch_model_ops.py

@@ -5,12 +5,14 @@ from . import ops
 
 LAYERS = {}
 
-@ops.op("Define PyTorch model", sub_nodes=LAYERS)
+op = ops.op_registration('LynxKite')
+
+@op("Define PyTorch model", sub_nodes=LAYERS)
 def define_pytorch_model(*, sub_flow):
   print('sub_flow:', sub_flow)
   return ops.Bundle(other={'model': str(sub_flow)})
 
-@ops.op("Train PyTorch model")
+@op("Train PyTorch model")
 def train_pytorch_model(model, graph):
   # import torch # Lazy import because it's slow.
   return 'hello ' + str(model)

server/workspace.py

@@ -37,48 +37,14 @@ class WorkspaceEdge(BaseConfig):
     target: str
 
 class Workspace(BaseConfig):
+    env: str = ''
     nodes: list[WorkspaceNode] = dataclasses.field(default_factory=list)
     edges: list[WorkspaceEdge] = dataclasses.field(default_factory=list)
 
 
-def execute(ws):
-    # Nodes are responsible for interpreting/executing their child nodes.
-    nodes = [n for n in ws.nodes if not n.parentId]
-    children = {}
-    for n in ws.nodes:
-        if n.parentId:
-            children.setdefault(n.parentId, []).append(n)
-    outputs = {}
-    failed = 0
-    while len(outputs) + failed < len(nodes):
-        for node in nodes:
-            if node.id in outputs:
-                continue
-            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]
-                data = node.data
-                op = ops.ALL_OPS[data.title]
-                params = {**data.params}
-                if op.sub_nodes:
-                    sub_nodes = children.get(node.id, [])
-                    sub_node_ids = [node.id for node in sub_nodes]
-                    sub_edges = [edge for edge in ws.edges if edge.source in sub_node_ids]
-                    params['sub_flow'] = {'nodes': sub_nodes, 'edges': sub_edges}
-                try:
-                  output = op(*inputs, **params)
-                except Exception as e:
-                  traceback.print_exc()
-                  data.error = str(e)
-                  failed += 1
-                  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)}
-                data.error = None
-                outputs[node.id] = output
-                if op.type == 'visualization' or op.type == 'table_view':
-                    data.view = output
+def execute(ws: Workspace):
+    if ws.env in ops.EXECUTORS:
+        ops.EXECUTORS[ws.env](ws)
 
 
 def save(ws: Workspace, path: str):
@@ -101,6 +67,7 @@ def load(path: str):
 
 
 def _update_metadata(ws):
+    catalog = ops.CATALOGS.get(ws.env, {})
     nodes = {node.id: node for node in ws.nodes}
     done = set()
     while len(done) < len(nodes):
@@ -109,7 +76,7 @@ def _update_metadata(ws):
                 continue
             data = node.data
             if node.parentId is None:
-                op = ops.ALL_OPS.get(data.title)
+                op = catalog.get(data.title)
             elif node.parentId not in nodes:
                 data.error = f'Parent not found: {node.parentId}'
                 done.add(node.id)

web/src/LynxKiteFlow.svelte

@@ -4,7 +4,6 @@
   import {
     SvelteFlow,
     Controls,
-    Background,
     MiniMap,
     MarkerType,
     useSvelteFlow,
@@ -14,6 +13,9 @@
     type Connection,
     type NodeTypes,
   } from '@xyflow/svelte';
+  import ArrowBack from 'virtual:icons/tabler/arrow-back'
+  import Backspace from 'virtual:icons/tabler/backspace'
+  import Atom from 'virtual:icons/tabler/Atom'
   import { useQuery, useMutation, useQueryClient } from '@sveltestack/svelte-query';
   import NodeWithParams from './NodeWithParams.svelte';
   import NodeWithVisualization from './NodeWithVisualization.svelte';
@@ -21,6 +23,7 @@
   import NodeWithSubFlow from './NodeWithSubFlow.svelte';
   import NodeWithArea from './NodeWithArea.svelte';
   import NodeSearch from './NodeSearch.svelte';
+  import EnvironmentSelector from './EnvironmentSelector.svelte';
   import '@xyflow/svelte/dist/style.css';
 
   export let path = '';
@@ -42,7 +45,7 @@
     return await res.json();
   }, {
     onSuccess: data => queryClient.setQueryData(['workspace', path], data),
-  })
+  });
 
   const nodeTypes: NodeTypes = {
     basic: NodeWithParams,
@@ -73,7 +76,7 @@
     event.preventDefault();
     nodeSearchSettings = {
       pos: { x: event.clientX, y: event.clientY },
-      boxes: $boxes,
+      boxes: $catalog.data[$backendWorkspace.data?.env],
     };
   }
   function addNode(e) {
@@ -106,13 +109,10 @@
     });
     closeNodeSearch();
   }
-  const boxes = writable([]);
-  async function getBoxes() {
+  const catalog = useQuery(['catalog'], async () => {
     const res = await fetch('/api/catalog');
-    const j = await res.json();
-    boxes.set(j);
-  }
-  getBoxes();
+    return res.json();
+  }, {staleTime: 60000, retry: false});
 
   let nodeSearchSettings: {
     pos: XYPosition,
@@ -153,6 +153,7 @@
       delete node.markerEnd;
       delete node.selected;
     }
+    g.env = $backendWorkspace.data?.env;
     const ws = orderedJSON(g);
     const bd = orderedJSON($backendWorkspace.data);
     if (ws === bd) return;
@@ -182,23 +183,74 @@
       parentId: node.id,
     };
   }
+  $: parentDir = path.split('/').slice(0, -1).join('/');
 
 </script>
 
-<div style:height="100%">
-  <SvelteFlow {nodes} {edges} {nodeTypes} fitView
-    on:paneclick={toggleNodeSearch}
-    on:nodeclick={nodeClick}
-    proOptions={{ hideAttribution: true }}
-    maxZoom={3}
-    minZoom={0.3}
-    onconnect={onconnect}
-    defaultEdgeOptions={{ markerEnd: { type: MarkerType.Arrow } }}
-    >
-    <Controls />
-    <MiniMap />
-    {#if nodeSearchSettings}
+<div class="page">
+  <div class="top-bar">
+    <div class="ws-name">
+      <a href><img src="/favicon.ico"></a>
+      {path}
+    </div>
+    <div class="tools">
+      <EnvironmentSelector
+        options={Object.keys($catalog.data || {})}
+        value={$backendWorkspace.data?.env}
+        onChange={(env) => $mutation.mutate({ path, ws: { ...$backendWorkspace.data, env } })}
+        />
+      <a href><Atom /></a>
+      <a href><Backspace /></a>
+      <a href="#dir?path={parentDir}"><ArrowBack /></a>
+    </div>
+  </div>
+  <div style:height="100%">
+    <SvelteFlow {nodes} {edges} {nodeTypes} fitView
+      on:paneclick={toggleNodeSearch}
+      on:nodeclick={nodeClick}
+      proOptions={{ hideAttribution: true }}
+      maxZoom={3}
+      minZoom={0.3}
+      onconnect={onconnect}
+      defaultEdgeOptions={{ markerEnd: { type: MarkerType.Arrow } }}
+      >
+      <Controls />
+      <MiniMap />
+      {#if nodeSearchSettings}
       <NodeSearch pos={nodeSearchSettings.pos} boxes={nodeSearchSettings.boxes} on:cancel={closeNodeSearch} on:add={addNode} />
-    {/if}
-  </SvelteFlow>
+      {/if}
+    </SvelteFlow>
+  </div>
 </div>
+
+<style>
+  .top-bar {
+    display: flex;
+    justify-content: space-between;
+    background: oklch(30% 0.13 230);
+    color: white;
+  }
+  .ws-name {
+    font-size: 1.5em;
+  }
+  .ws-name img {
+    height: 1.5em;
+    vertical-align: middle;
+    margin: 4px;
+  }
+  .page {
+    display: flex;
+    flex-direction: column;
+    height: 100vh;
+  }
+
+  .tools {
+    display: flex;
+    align-items: center;
+  }
+  .tools a {
+    color: oklch(75% 0.13 230);
+    font-size: 1.5em;
+    padding: 0 10px;
+  }
+</style>

web/src/Workspace.svelte

@@ -2,62 +2,13 @@
   // This is the whole LynxKite workspace editor page.
   import { QueryClient, QueryClientProvider } from '@sveltestack/svelte-query'
   import { SvelteFlowProvider } from '@xyflow/svelte';
-  import ArrowBack from 'virtual:icons/tabler/arrow-back'
-  import Backspace from 'virtual:icons/tabler/backspace'
-  import Atom from 'virtual:icons/tabler/Atom'
   import LynxKiteFlow from './LynxKiteFlow.svelte';
   export let path = '';
-  $: parent = path.split('/').slice(0, -1).join('/');
   const queryClient = new QueryClient()
 </script>
 
 <QueryClientProvider client={queryClient}>
-  <div class="page">
-    <div class="top-bar">
-      <div class="ws-name">
-        <a href><img src="/favicon.ico"></a>
-        {path}
-      </div>
-      <div class="tools">
-        <a href><Atom /></a>
-        <a href><Backspace /></a>
-        <a href="#dir?path={parent}"><ArrowBack /></a>
-      </div>
-    </div>
-    <SvelteFlowProvider>
-      <LynxKiteFlow path={path} />
-    </SvelteFlowProvider>
-  </div>
+  <SvelteFlowProvider>
+    <LynxKiteFlow path={path} />
+  </SvelteFlowProvider>
 </QueryClientProvider>
-
-<style>
-  .top-bar {
-    display: flex;
-    justify-content: space-between;
-    background: oklch(30% 0.13 230);
-    color: white;
-  }
-  .ws-name {
-    font-size: 1.5em;
-  }
-  .ws-name img {
-    height: 1.5em;
-    vertical-align: middle;
-    margin: 4px;
-  }
-  .page {
-    display: flex;
-    flex-direction: column;
-    height: 100vh;
-  }
-
-  .tools {
-    display: flex;
-    align-items: center;
-  }
-  .tools a {
-    color: oklch(75% 0.13 230);
-    font-size: 1.5em;
-    padding: 0 10px;
-  }
-</style>