Start writing "LLM logic" executor.

server/llm_ops.py

@@ -1,13 +1,29 @@
 '''For specifying an LLM agent logic flow.'''
 from . import ops
+import dataclasses
+import inspect
 import json
 import openai
 import pandas as pd
+import traceback
+from . import workspace
 
 client = openai.OpenAI(base_url="http://localhost:11434/v1")
 CACHE = {}
+ENV = 'LLM logic'
+op = ops.op_registration(ENV)
 
-op = ops.op_registration('LLM logic')
+@dataclasses.dataclass
+class Context:
+  '''Passed to operation functions as "_ctx" if they have such a parameter.'''
+  node: workspace.WorkspaceNode
+  last_result = None
+
+@dataclasses.dataclass
+class Output:
+  '''Return this to send values to specific outputs of a node.'''
+  output_handle: str
+  value: dict
 
 def chat(*args, **kwargs):
   key = json.dumps({'args': args, 'kwargs': kwargs})
@@ -23,52 +39,129 @@ def input(*, filename: ops.PathStr, key: str):
 @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()
-  prompts = []
-  for i, row in df.iterrows():
-    p = template
-    for c in df.columns:
-      p = p.replace(c.upper(), str(row[c]))
-    prompts.append(p)
-  df['prompt'] = prompts
-  return df
-
+  p = template
+  for k, v in input.items():
+    p = p.replace(k.upper(), str(v))
+  return p
 
 @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 = {}
+  assert 'prompt' in input, 'Please create the prompt first.'
+  options = {}
   if accepted_regex:
-    g['extra_body'] = {
+    options['extra_body'] = {
       "guided_regex": accepted_regex,
     }
-  for i, row in df.iterrows():
-    [res] = chat(
-      model=model,
-      max_tokens=max_tokens,
-      messages=[
-        {"role": "user", "content": row['prompt']},
-      ],
-      **g,
-    )
-    df.loc[i, 'response'] = res
-  return df
+  results = chat(
+    model=model,
+    max_tokens=max_tokens,
+    messages=[
+      {"role": "user", "content": input['prompt']},
+    ],
+    **options,
+  )
+  return [{**input, 'response': r} for r in results]
 
 @op("View", view="table_view")
-def view(input):
-  v = {
-    'dataframes': { 'df': {
-      'columns': [str(c) for c in input.columns],
-      'data': input.values.tolist(),
-    }}
-  }
+def view(input, *, _ctx: Context):
+  v = _ctx.last_result
+  if v:
+    columns = v['dataframes']['df']['columns']
+    v['dataframes']['df']['data'].append([input[c] for c in columns])
+  else:
+    columns = [str(c) for c in input.keys() if not str(c).startswith('_')]
+    v = {
+      'dataframes': { 'df': {
+        'columns': columns,
+        'data': [input[c] for c in columns],
+      }}
+    }
   return v
 
 @ops.input_position(input="right")
 @ops.output_position(output="left")
 @op("Loop")
-def loop(input, *, max_iterations: int = 10):
-  '''Data can flow back here until it becomes empty or reaches the limit.'''
+def loop(input, *, max_iterations: int = 3, _ctx: Context):
+  '''Data can flow back here max_iterations-1 times.'''
+  key = f'iterations-{_ctx.node.id}'
+  input[key] = input.get(key, 0) + 1
+  if input[key] < max_iterations:
+    return input
+
+@op('Branch', outputs=['true', 'false'])
+def branch(input, *, expression: str):
+  res = eval(expression, input)
+  return Output(str(bool(res)).lower(), input)
+
+@ops.input_position(db="top")
+@op('RAG')
+def rag(input, db, *, closest_n: int=10):
   return input
+
+@op('Run Python')
+def run_python(input, *, template: str):
+  assert template, 'Please specify the template. Refer to columns using their names in uppercase.'
+  p = template
+  for k, v in input.items():
+    p = p.replace(k.upper(), str(v))
+  return p
+
+
+
+@ops.register_executor(ENV)
+def execute(ws):
+  catalog = ops.CATALOGS[ENV]
+  nodes = {n.id: n for n in ws.nodes}
+  contexts = {n.id: Context(n) for n in ws.nodes}
+  edges = {n.id: [] for n in ws.nodes}
+  for e in ws.edges:
+    edges[e.source].append(e.target)
+  tasks = {}
+  NO_INPUT = object() # Marker for initial tasks.
+  for node in ws.nodes:
+    node.data.error = None
+    op = catalog[node.data.title]
+    # Start tasks for nodes that have no inputs.
+    if not op.inputs:
+      tasks[node.id] = [NO_INPUT]
+  # Run the rest until we run out of tasks.
+  while tasks:
+    n, ts = tasks.popitem()
+    node = nodes[n]
+    data = node.data
+    op = catalog[data.title]
+    params = {**data.params}
+    if has_ctx(op):
+      params['_ctx'] = contexts[node.id]
+    results = []
+    for task in ts:
+      try:
+        if task is NO_INPUT:
+          result = op(**params)
+        else:
+          # TODO: Tasks with multiple inputs?
+          result = op(task, **params)
+      except Exception as e:
+        traceback.print_exc()
+        data.error = str(e)
+        break
+      contexts[node.id].last_result = result
+      # 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)
+    else: # Finished all tasks without errors.
+      if op.type == 'visualization' or op.type == 'table_view':
+        data.display = results
+      for target in edges[node.id]:
+        tasks.setdefault(target, []).extend(results)
+
+def df_to_list(df):
+  return [dict(zip(df.columns, row)) for row in df.values]
+
+def has_ctx(op):
+  sig = inspect.signature(op.func)
+  return '_ctx' in sig.parameters

server/lynxkite_ops.py

@@ -1,12 +1,135 @@
 '''Some operations. To be split into separate files when we have more.'''
 from . import ops
+import dataclasses
+import functools
 import matplotlib
 import networkx as nx
 import pandas as pd
 import traceback
+import typing
 
 op = ops.op_registration('LynxKite')
 
+@dataclasses.dataclass
+class RelationDefinition:
+  '''Defines a set of edges.'''
+  df: str # The DataFrame that contains the edges.
+  source_column: str # The column in the edge DataFrame that contains the source node ID.
+  target_column: str # The column in the edge DataFrame that contains the target node ID.
+  source_table: str # The DataFrame that contains the source nodes.
+  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.
+
+@dataclasses.dataclass
+class Bundle:
+  '''A collection of DataFrames and other data.
+
+  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[str, pd.DataFrame] = dataclasses.field(default_factory=dict)
+  relations: list[RelationDefinition] = dataclasses.field(default_factory=list)
+  other: dict[str, typing.Any] = None
+
+  @classmethod
+  def from_nx(cls, graph: nx.Graph):
+    edges = nx.to_pandas_edgelist(graph)
+    d = dict(graph.nodes(data=True))
+    nodes = pd.DataFrame(d.values(), index=d.keys())
+    nodes['id'] = nodes.index
+    return cls(
+      dfs={'edges': edges, 'nodes': nodes},
+      relations=[
+        RelationDefinition(
+          df='edges',
+          source_column='source',
+          target_column='target',
+          source_table='nodes',
+          target_table='nodes',
+          source_key='id',
+          target_key='id',
+        )
+      ]
+    )
+
+  def to_nx(self):
+    graph = nx.from_pandas_edgelist(self.dfs['edges'])
+    nx.set_node_attributes(graph, self.dfs['nodes'].set_index('id').to_dict('index'))
+    return graph
+
+
+def nx_node_attribute_func(name):
+  '''Decorator for wrapping a function that adds a NetworkX node attribute.'''
+  def decorator(func):
+    @functools.wraps(func)
+    def wrapper(graph: nx.Graph, **kwargs):
+      graph = graph.copy()
+      attr = func(graph, **kwargs)
+      nx.set_node_attributes(graph, attr, name)
+      return graph
+    return wrapper
+  return decorator
+
+
+def disambiguate_edges(ws):
+    '''If an input plug is connected to multiple edges, keep only the last edge.'''
+    seen = set()
+    for edge in reversed(ws.edges):
+        if (edge.target, edge.targetHandle) in seen:
+            ws.edges.remove(edge)
+        seen.add((edge.target, edge.targetHandle))
+
+
+@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]
+    disambiguate_edges(ws)
+    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
+            # 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]
+                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}
+                # Convert inputs.
+                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)
+                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.display = output
+
 @op("Import Parquet")
 def import_parquet(*, filename: str):
   '''Imports a parquet file.'''
@@ -18,7 +141,7 @@ def create_scale_free_graph(*, nodes: int = 10):
   return nx.scale_free_graph(nodes)
 
 @op("Compute PageRank")
-@ops.nx_node_attribute_func('pagerank')
+@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)
 
@@ -26,7 +149,7 @@ def compute_pagerank(graph: nx.Graph, *, damping=0.85, iterations=100):
 @op("Sample graph")
 def sample_graph(graph: nx.Graph, *, nodes: int = 100):
   '''Takes a subgraph.'''
-  return nx.scale_free_graph(nodes)
+  return nx.scale_free_graph(nodes) # TODO: Implement this.
 
 
 def _map_color(value):
@@ -36,7 +159,7 @@ def _map_color(value):
   return ['#{:02x}{:02x}{:02x}'.format(int(r*255), int(g*255), int(b*255)) for r, g, b in rgba[:, :3]]
 
 @op("Visualize graph", view="visualization")
-def visualize_graph(graph: ops.Bundle, *, color_nodes_by: 'node_attribute' = None):
+def visualize_graph(graph: Bundle, *, color_nodes_by: 'node_attribute' = None):
   nodes = graph.dfs['nodes'].copy()
   if color_nodes_by:
     nodes['color'] = _map_color(nodes[color_nodes_by])
@@ -79,7 +202,7 @@ def visualize_graph(graph: ops.Bundle, *, color_nodes_by: 'node_attribute' = Non
   return v
 
 @op("View tables", view="table_view")
-def view_tables(bundle: ops.Bundle):
+def view_tables(bundle: Bundle):
   v = {
     'dataframes': { name: {
       'columns': [str(c) for c in df.columns],
@@ -89,44 +212,3 @@ 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/ops.py

@@ -95,80 +95,12 @@ class Op(BaseConfig):
           params[p] = int(params[p])
         elif self.params[p].type == float:
           params[p] = float(params[p])
-    # Convert inputs.
-    inputs = list(inputs)
-    for i, (x, p) in enumerate(zip(inputs, self.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)
     res = self.func(*inputs, **params)
     return res
 
 
-@dataclasses.dataclass
-class RelationDefinition:
-  '''Defines a set of edges.'''
-  df: str # The DataFrame that contains the edges.
-  source_column: str # The column in the edge DataFrame that contains the source node ID.
-  target_column: str # The column in the edge DataFrame that contains the target node ID.
-  source_table: str # The DataFrame that contains the source nodes.
-  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.
-
-@dataclasses.dataclass
-class Bundle:
-  '''A collection of DataFrames and other data.
-
-  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[str, pd.DataFrame] = dataclasses.field(default_factory=dict)
-  relations: list[RelationDefinition] = dataclasses.field(default_factory=list)
-  other: dict[str, typing.Any] = None
-
-  @classmethod
-  def from_nx(cls, graph: nx.Graph):
-    edges = nx.to_pandas_edgelist(graph)
-    d = dict(graph.nodes(data=True))
-    nodes = pd.DataFrame(d.values(), index=d.keys())
-    nodes['id'] = nodes.index
-    return cls(
-      dfs={'edges': edges, 'nodes': nodes},
-      relations=[
-        RelationDefinition(
-          df='edges',
-          source_column='source',
-          target_column='target',
-          source_table='nodes',
-          target_table='nodes',
-          source_key='id',
-          target_key='id',
-        )
-      ]
-    )
-
-  def to_nx(self):
-    graph = nx.from_pandas_edgelist(self.dfs['edges'])
-    nx.set_node_attributes(graph, self.dfs['nodes'].set_index('id').to_dict('index'))
-    return graph
-
-
-def nx_node_attribute_func(name):
-  '''Decorator for wrapping a function that adds a NetworkX node attribute.'''
-  def decorator(func):
-    @functools.wraps(func)
-    def wrapper(graph: nx.Graph, **kwargs):
-      graph = graph.copy()
-      attr = func(graph, **kwargs)
-      nx.set_node_attributes(graph, attr, name)
-      return graph
-    return wrapper
-  return decorator
-
 
-def op(env: str, name: str, *, view='basic', sub_nodes=None):
+def op(env: str, name: str, *, view='basic', sub_nodes=None, outputs=None):
   '''Decorator for defining an operation.'''
   def decorator(func):
     sig = inspect.signature(func)
@@ -179,10 +111,13 @@ def op(env: str, name: str, *, view='basic', sub_nodes=None):
       if param.kind != param.KEYWORD_ONLY}
     params = {}
     for n, param in sig.parameters.items():
-      if param.kind == param.KEYWORD_ONLY:
+      if param.kind == param.KEYWORD_ONLY and not n.startswith('_'):
         params[n] = Parameter.basic(n, param.default, param.annotation)
-    outputs = {'output': Output(name='output', type=None)} if view == 'basic' else {} # Maybe more fancy later.
-    op = Op(func=func, name=name, params=params, inputs=inputs, outputs=outputs, type=view)
+    if outputs:
+      _outputs = {name: Output(name=name, type=None) for name in outputs}
+    else:
+      _outputs = {'output': Output(name='output', type=None)} if view == 'basic' else {}
+    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'
@@ -213,7 +148,7 @@ def output_position(**kwargs):
 def no_op(*args, **kwargs):
   if args:
     return args[0]
-  return Bundle()
+  return None
 
 def register_passive_op(env: str, name: str, inputs=[], outputs=['output'], params=[]):
   '''A passive operation has no associated code.'''

server/test_llm_ops.py

@@ -0,0 +1,28 @@
+import unittest
+from . import llm_ops
+from . import workspace
+
+class LLMOpsTest(unittest.TestCase):
+  def testExecute(self):
+    ws = workspace.Workspace(env='LLM logic', nodes=[
+      workspace.WorkspaceNode(
+        id='0',
+        type='basic',
+        position=workspace.Position(x=0, y=0),
+        data=workspace.WorkspaceNodeData(title='Input', params={
+          'filename': '/Users/danieldarabos/Downloads/aimo-train.csv',
+          'key': 'problem',
+        })),
+      workspace.WorkspaceNode(
+        id='1',
+        type='table_view',
+        position=workspace.Position(x=0, y=0),
+        data=workspace.WorkspaceNodeData(title='View', params={})),
+    ], edges=[
+      workspace.WorkspaceEdge(id='0-1', source='0', target='1', sourceHandle='', targetHandle=''),
+    ])
+    llm_ops.execute(ws)
+    self.assertEqual('', ws.nodes[1].data.display)
+
+if __name__ == '__main__':
+  unittest.main()

server/workspace.py

@@ -35,6 +35,8 @@ class WorkspaceEdge(BaseConfig):
     id: str
     source: str
     target: str
+    sourceHandle: str
+    targetHandle: str
 
 class Workspace(BaseConfig):
     env: str = ''

web/src/Directory.svelte

@@ -58,7 +58,9 @@
   </div>
   <div class="entry-list">
     {#await list}
-      <div>Loading...</div>
+      <div class="loading spinner-border" role="status">
+        <span class="visually-hidden">Loading...</span>
+      </div>
     {:then list}
       <div class="actions">
         <a href="{newWorkspaceIn(path, list)}"><FilePlus /> New workspace</a>
@@ -168,4 +170,8 @@
     color: black;
     text-decoration: none;
   }
+  .loading {
+    color: #39bcf3;
+    margin: 10px;
+  }
 </style>

web/src/EnvironmentSelector.svelte

@@ -0,0 +1,14 @@
+<script lang="ts">
+  export let options;
+  export let value;
+  export let onChange;
+</script>
+
+<select class="form-select form-select-sm"
+  value={value}
+  on:change={(evt) => onChange(evt.currentTarget.value)}
+  >
+  {#each options as option}
+    <option value={option}>{option}</option>
+  {/each}
+</select>

web/src/LynxKiteFlow.svelte

@@ -160,15 +160,6 @@
     console.log('changed', JSON.stringify(diff(g, $backendWorkspace.data), null, 2));
     $mutation.mutate({ path, ws: g });
   }
-  function onconnect(connection: Connection) {
-    edges.update((edges) => {
-      // Only one source can connect to a given target.
-      return edges.filter((e) =>
-        e.source === connection.source
-        || e.target !== connection.target
-        || e.targetHandle !== connection.targetHandle);
-    });
-  }
   function nodeClick(e) {
     const node = e.detail.node;
     const meta = node.data.meta;
@@ -211,7 +202,6 @@
       proOptions={{ hideAttribution: true }}
       maxZoom={3}
       minZoom={0.3}
-      onconnect={onconnect}
       defaultEdgeOptions={{ markerEnd: { type: MarkerType.Arrow } }}
       >
       <Controls />

web/src/NodeWithTableView.svelte

@@ -6,18 +6,18 @@
   type $$Props = NodeProps;
   export let data: $$Props['data'];
   const open = {};
-  $: single = data.view?.dataframes && Object.keys(data.view.dataframes).length === 1;
+  $: single = data.display?.dataframes && Object.keys(data.display.dataframes).length === 1;
 </script>
 
 <LynxKiteNode {...$$props}>
-  {#if data.view}
-    {#each Object.entries(data.view.dataframes) as [name, df]}
+  {#if data.display}
+    {#each Object.entries(data.display.dataframes) as [name, df]}
       {#if !single}<div class="df-head" on:click={() => open[name] = !open[name]}>{name}</div>{/if}
       {#if single || open[name]}
         <Table columns={df.columns} data={df.data} />
       {/if}
     {/each}
-    {#each Object.entries(data.view.others || {}) as [name, o]}
+    {#each Object.entries(data.display.others || {}) as [name, o]}
       <div class="df-head" on:click={() => open[name] = !open[name]}>{name}</div>
       {#if open[name]}
       <pre>{o}</pre>

web/src/NodeWithVisualization.svelte

@@ -8,8 +8,8 @@
 </script>
 
 <NodeWithParams {...$$props}>
-  {#if data.view}
-    <Chart {init} options={data.view} initOptions={{renderer: 'canvas', width: 250, height: 250}}/>
+  {#if data.display}
+    <Chart {init} options={data.display} initOptions={{renderer: 'canvas', width: 250, height: 250}}/>
   {/if}
 </NodeWithParams>
 <style>