Spaces:
Running
Running
File size: 2,580 Bytes
a07e9cb eda8f97 a07e9cb eda8f97 a07e9cb 6988728 eda8f97 6988728 a07e9cb 6988728 eda8f97 6988728 a07e9cb 6988728 a07e9cb eda8f97 a07e9cb 6988728 eda8f97 6988728 a07e9cb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 |
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()
|