server/test_llm_ops.py

import unittest
from . import ops
from . import llm_ops
from .executors import one_by_one
from . import workspace

def make_node(id, op, type='basic', **params):
  return workspace.WorkspaceNode(
    id=id,
    type=type,
    position=workspace.Position(x=0, y=0),
    data=workspace.WorkspaceNodeData(title=op, params=params),
  )
def make_input(id):
  return make_node(
    id, 'Input CSV',
    filename='/Users/danieldarabos/Downloads/aimo-train.csv',
    key='problem')
def make_edge(source, target, targetHandle='input'):
  return workspace.WorkspaceEdge(
    id=f'{source}-{target}', source=source, target=target, sourceHandle='', targetHandle=targetHandle)

class LLMOpsTest(unittest.TestCase):
  def testExecute(self):
    ws = workspace.Workspace(env='LLM logic', nodes=[
      make_node(
        '0', 'Input CSV',
        filename='/Users/danieldarabos/Downloads/aimo-train.csv',
        key='problem'),
      make_node(
        '1', 'View', type='table_view'),
    ], edges=[
      make_edge('0', '1')
    ])
    catalog = ops.CATALOGS[ws.env]
    one_by_one.execute(ws, catalog)
    # self.assertEqual('', ws.nodes[1].data.display)

  def testStages(self):
    ws = workspace.Workspace(env='LLM logic', nodes=[
      make_input('in1'), make_input('in2'), make_input('in3'),
      make_node('rag1', 'RAG'), make_node('rag2', 'RAG'),
      make_node('p1', 'Create prompt'), make_node('p2', 'Create prompt'),
    ], edges=[
      make_edge('in1', 'rag1', 'db'), make_edge('in2', 'rag1'),
      make_edge('rag1', 'p1'), make_edge('p1', 'rag2', 'db'),
      make_edge('in3', 'p2'), make_edge('p3', 'rag2'),
    ])
    catalog = ops.CATALOGS[ws.env]
    stages = one_by_one.get_stages(ws, catalog)
    print(stages)
    # self.assertEqual('', stages)

  def testStagesMultiInput(self):
    ws = workspace.Workspace(env='LLM logic', nodes=[
      make_node('doc', 'Input document'),
      make_node('split', 'Split document'),
      make_node('graph', 'Build document graph'),
      make_node('chat', 'Input chat'),
      make_node('rag', 'RAG'),
      make_node('neighbors', 'Add neighbors'),
    ], edges=[
      make_edge('doc', 'split'), make_edge('split', 'graph'),
      make_edge('split', 'rag', 'db'), make_edge('chat', 'rag', 'input'),
      make_edge('split', 'neighbors', 'nodes'),
      make_edge('graph', 'neighbors', 'edges'),
      make_edge('rag', 'neighbors', 'item'),
    ])
    catalog = ops.CATALOGS[ws.env]
    stages = one_by_one.get_stages(ws, catalog)
    print(stages)
    # self.assertEqual('', stages)

if __name__ == '__main__':
  unittest.main()