From 69e7ec01cb2422383e07510aececcd4009048d97 Mon Sep 17 00:00:00 2001 From: zhangyongchao Date: Tue, 12 Nov 2024 12:01:51 +0800 Subject: [PATCH 01/12] fix bug: map backend_nodeid to frontend_nodeid --- lazyllm/engine/engine.py | 11 +- lazyllm/engine/node.py | 3 +- tests/basic_tests/test_engine.py | 455 ++++++++++++++++++++++++++----- 3 files changed, 399 insertions(+), 70 deletions(-) diff --git a/lazyllm/engine/engine.py b/lazyllm/engine/engine.py index eb4677af..1a3e54fe 100644 --- a/lazyllm/engine/engine.py +++ b/lazyllm/engine/engine.py @@ -131,10 +131,15 @@ def get_args(cls, key, value, builder_key=None): return node def _process_hook(self, node, module): - if not isinstance(module, (lazyllm.ModuleBase, lazyllm.LazyLLMFlowsBase)): + if not node.enable_data_reflow: return - if node.enable_data_reflow: - node.func.register_hook(NodeMetaHook) + if isinstance(module, lazyllm.ModuleBase): + NodeMetaHook.MODULEID_TO_WIDGETID[module._module_id] = node.id + elif isinstance(module, lazyllm.LazyLLMFlowsBase): + NodeMetaHook.MODULEID_TO_WIDGETID[module._flow_id] = node.id + else: + return + node.func.register_hook(NodeMetaHook) _constructor = NodeConstructor() diff --git a/lazyllm/engine/node.py b/lazyllm/engine/node.py index 4615f3d3..0ab31eae 100644 --- a/lazyllm/engine/node.py +++ b/lazyllm/engine/node.py @@ -9,7 +9,7 @@ @dataclass class Node(): - id: int + id: str kind: str name: str args: Optional[Dict] = None @@ -17,7 +17,6 @@ class Node(): func: Optional[Callable] = None arg_names: Optional[List[str]] = None enable_data_reflow: bool = False - enable_hook: bool = False subitem_name: Optional[Union[List[str], str]] = None @property diff --git a/tests/basic_tests/test_engine.py b/tests/basic_tests/test_engine.py index b33f4cd6..cfab0269 100644 --- a/tests/basic_tests/test_engine.py +++ b/tests/basic_tests/test_engine.py @@ -60,7 +60,14 @@ def test_engine_subgraph(self): assert '1234' in r def test_engine_code(self): - nodes = [dict(id='1', kind='Code', name='m1', args='def test(x: int):\n return 2 * x\n')] + nodes = [ + dict( + id="1", + kind="Code", + name="m1", + args=dict(code="def test(x: int):\n return 2 * x\n"), + ) + ] edges = [dict(iid='__start__', oid='1'), dict(iid='1', oid='__end__')] engine = LightEngine() @@ -69,9 +76,28 @@ def test_engine_code(self): assert engine.run(gid, 2) == 4 def test_engine_switch(self): - plus1 = dict(id='1', kind='Code', name='m1', args='def test(x: int):\n return 1 + x\n') - double = dict(id='2', kind='Code', name='m2', args='def test(x: int):\n return 2 * x\n') - square = dict(id='3', kind='Code', name='m3', args='def test(x: int):\n return x * x\n') + plus1 = dict( + id="1", + kind="Code", + name="m1", + args=dict(code="def test(x: int):\n return 1 + x\n"), + ) + double = dict( + id="2", + kind="Code", + name="m2", + args=dict(code="def test(x: int):\n return 2 * x\n"), + ) + # square = dict(id='3', kind='Code', name='m3', args='def test(x: int):\n return x * x\n') + square = dict( + id="3", + kind="Code", + name="m3", + args=dict( + code="def test(x: int):\n return x * x\n", + _lazyllm_enable_report=True, + ), + ) switch = dict( id="4", kind="Switch", @@ -92,6 +118,7 @@ def test_engine_switch(self): assert engine.run(gid, 3) == 9 engine.reset() + lazyllm.globals._init_sid("session_test_002") switch = dict( id="4", @@ -112,9 +139,28 @@ def test_engine_switch(self): assert engine.run(gid, 'case3', 3) == 9 def test_engine_ifs(self): - plus1 = dict(id='1', kind='Code', name='m1', args='def test(x: int):\n return 1 + x\n') - double = dict(id='2', kind='Code', name='m2', args='def test(x: int):\n return 2 * x\n') - square = dict(id='3', kind='Code', name='m3', args='def test(x: int):\n return x * x\n') + # plus1 = dict(id='1', kind='Code', name='m1', args='def test(x: int):\n return 1 + x\n') + plus1 = dict( + id="1", + kind="Code", + name="m1", + args=dict( + code="def test(x: int):\n return 1 + x\n", + _lazyllm_enable_report=True, + ), + ) + double = dict( + id="2", + kind="Code", + name="m2", + args=dict(code="def test(x: int):\n return 2 * x\n"), + ) + square = dict( + id="3", + kind="Code", + name="m3", + args=dict(code="def test(x: int):\n return x * x\n"), + ) ifs = dict( id="4", kind="Ifs", @@ -136,7 +182,14 @@ def test_engine_ifs(self): assert engine.run(gid, 10) == 100 def test_engine_loop(self): - nodes = [dict(id='1', kind='Code', name='code', args='def square(x: int): return x * x')] + nodes = [ + dict( + id="1", + kind="Code", + name="code", + args=dict(code="def square(x: int): return x * x"), + ) + ] edges = [dict(iid='__start__', oid='1'), dict(iid='1', oid='__end__')] nodes = [ @@ -160,7 +213,14 @@ def test_engine_loop(self): assert engine.run(gid, 2) == 16 def test_engine_warp(self): - nodes = [dict(id='1', kind='Code', name='code', args='def square(x: int): return x * x')] + nodes = [ + dict( + id="1", + kind="Code", + name="code", + args=dict(code="def square(x: int): return x * x"), + ) + ] edges = [dict(iid='__start__', oid='1'), dict(iid='1', oid='__end__')] nodes = [ @@ -222,10 +282,34 @@ def test_engine_formatter(self): '{"query": "hiaha", "files": ["path/to/file"]}' def test_engine_edge_formatter(self): - nodes = [dict(id='1', kind='Code', name='m1', args='def test(x: int):\n return x\n'), - dict(id='2', kind='Code', name='m2', args='def test(x: int):\n return [[x, 2*x], [3*x, 4*x]]\n'), - dict(id='3', kind='Code', name='m3', args='def test(x: int):\n return dict(a=1, b=x * x)\n'), - dict(id='4', kind='Code', name='m4', args='def test(x, y, z):\n return f"{x}{y}{z}"\n')] + nodes = [ + dict( + id="1", + kind="Code", + name="m1", + args=dict(code="def test(x: int):\n return x\n"), + ), + dict( + id="2", + kind="Code", + name="m2", + args=dict( + code="def test(x: int):\n return [[x, 2*x], [3*x, 4*x]]\n" + ), + ), + dict( + id="3", + kind="Code", + name="m3", + args=dict(code="def test(x: int):\n return dict(a=1, b=x * x)\n"), + ), + dict( + id="4", + kind="Code", + name="m4", + args=dict(code='def test(x, y, z):\n return f"{x}{y}{z}"\n'), + ), + ] edges = [dict(iid='__start__', oid='1'), dict(iid='__start__', oid='2'), dict(iid='__start__', oid='3'), dict(iid='1', oid='4'), dict(iid='2', oid='4', formatter='[:][1]'), dict(iid='3', oid='4', formatter='[b]'), dict(iid='4', oid='__end__')] @@ -236,9 +320,26 @@ def test_engine_edge_formatter(self): assert engine.run(gid, 2) == '2[4, 8]4' def test_engine_edge_formatter_start(self): - nodes = [dict(id='1', kind='Code', name='m1', args='def test(x: int): return x'), - dict(id='2', kind='Code', name='m2', args='def test(x: int): return 2 * x'), - dict(id='3', kind='Code', name='m3', args='def test(x, y): return x + y')] + nodes = [ + dict( + id="1", + kind="Code", + name="m1", + args=dict(code="def test(x: int): return x"), + ), + dict( + id="2", + kind="Code", + name="m2", + args=dict(code="def test(x: int): return 2 * x"), + ), + dict( + id="3", + kind="Code", + name="m3", + args=dict(code="def test(x, y): return x + y"), + ), + ] edges = [dict(iid='__start__', oid='1', formatter='[0]'), dict(iid='__start__', oid='2', formatter='[1]'), dict(iid='1', oid='3'), dict(iid='2', oid='3'), dict(iid='3', oid='__end__')] @@ -248,11 +349,35 @@ def test_engine_edge_formatter_start(self): assert engine.run(gid, 5, 3, 1) == 11 def test_engine_formatter_end(self): - nodes = [dict(id='1', kind='Code', name='m1', args='def test(x: int):\n return x\n'), - dict(id='2', kind='Code', name='m2', args='def test1(x: int):\n return [[x, 2*x], [3*x, 4*x]]\n'), - # two unused node - dict(id='3', kind='Code', name='m3', args='def test2(x: int):\n return dict(a=1, b=x * x)\n'), - dict(id='4', kind='Code', name='m4', args='def test3(x, y, z):\n return f"{x}{y}{z}"\n')] + nodes = [ + dict( + id="1", + kind="Code", + name="m1", + args=dict(code="def test(x: int):\n return x\n"), + ), + dict( + id="2", + kind="Code", + name="m2", + args=dict( + code="def test1(x: int):\n return [[x, 2*x], [3*x, 4*x]]\n" + ), + ), + # two unused node + dict( + id="3", + kind="Code", + name="m3", + args=dict(code="def test2(x: int):\n return dict(a=1, b=x * x)\n"), + ), + dict( + id="4", + kind="Code", + name="m4", + args=dict(code='def test3(x, y, z):\n return f"{x}{y}{z}"\n'), + ), + ] edges = [dict(iid='__start__', oid='1'), dict(iid='__start__', oid='2'), dict(iid='2', oid='__end__'), dict(iid='1', oid='__end__')] @@ -264,9 +389,28 @@ def test_engine_formatter_end(self): engine.reset() - nodes = [dict(id='1', kind='Code', name='m1', args='def test(x: int):\n return x\n'), - dict(id='2', kind='Code', name='m2', args='def test1(x: int):\n return [[x, 2*x], [3*x, 4*x]]\n'), - dict(id='3', kind='JoinFormatter', name='join', args=dict(type='to_dict', names=['a', 'b']))] + nodes = [ + dict( + id="1", + kind="Code", + name="m1", + args=dict(code="def test(x: int):\n return x\n"), + ), + dict( + id="2", + kind="Code", + name="m2", + args=dict( + code="def test1(x: int):\n return [[x, 2*x], [3*x, 4*x]]\n" + ), + ), + dict( + id="3", + kind="JoinFormatter", + name="join", + args=dict(type="to_dict", names=["a", "b"]), + ), + ] edges = [dict(iid='__start__', oid='1'), dict(iid='__start__', oid='2'), dict(iid='2', oid='3'), dict(iid='1', oid='3'), dict(iid='3', oid='__end__', formatter='*[a, b]')] engine = LightEngine() @@ -276,8 +420,15 @@ def test_engine_formatter_end(self): print(isinstance(r, lazyllm.package)) def test_engine_join_stack(self): - nodes = [dict(id='0', kind='Code', name='c1', args='def test(x: int): return x'), - dict(id='1', kind='JoinFormatter', name='join', args=dict(type='stack'))] + nodes = [ + dict( + id="0", + kind="Code", + name="c1", + args=dict(code="def test(x: int): return x"), + ), + dict(id="1", kind="JoinFormatter", name="join", args=dict(type="stack")), + ] edges = [dict(iid='__start__', oid='0'), dict(iid='0', oid='1'), dict(iid='1', oid='__end__')] engine = LightEngine() gid = engine.start(nodes, edges) @@ -287,10 +438,27 @@ def test_engine_join_stack(self): engine.reset() - nodes = [dict(id='0', kind='Code', name='c1', args='def test(x: int): return x'), - dict(id='1', kind='Code', name='c2', args='def test(x: int): return 2 * x'), - dict(id='2', kind='Code', name='c3', args='def test(x: int): return 3 * x'), - dict(id='3', kind='JoinFormatter', name='join', args=dict(type='stack'))] + nodes = [ + dict( + id="0", + kind="Code", + name="c1", + args=dict(code="def test(x: int): return x"), + ), + dict( + id="1", + kind="Code", + name="c2", + args=dict(code="def test(x: int): return 2 * x"), + ), + dict( + id="2", + kind="Code", + name="c3", + args=dict(code="def test(x: int): return 3 * x"), + ), + dict(id="3", kind="JoinFormatter", name="join", args=dict(type="stack")), + ] edges = [dict(iid='__start__', oid='0'), dict(iid='__start__', oid='1'), dict(iid='__start__', oid='2'), dict(iid='0', oid='3'), dict(iid='1', oid='3'), dict(iid='2', oid='3'), dict(iid='3', oid='__end__')] gid = engine.start(nodes, edges) @@ -299,8 +467,15 @@ def test_engine_join_stack(self): assert engine.run(gid, [1]) == [[1], [1, 1], [1, 1, 1]] def test_engine_join_sum(self): - nodes = [dict(id='0', kind='Code', name='c1', args='def test(x: int): return [x, 2 * x]'), - dict(id='1', kind='JoinFormatter', name='join', args=dict(type='sum'))] + nodes = [ + dict( + id="0", + kind="Code", + name="c1", + args=dict(code="def test(x: int): return [x, 2 * x]"), + ), + dict(id="1", kind="JoinFormatter", name="join", args=dict(type="sum")), + ] edges = [dict(iid='__start__', oid='0'), dict(iid='0', oid='1'), dict(iid='1', oid='__end__')] engine = LightEngine() gid = engine.start(nodes, edges) @@ -310,10 +485,27 @@ def test_engine_join_sum(self): engine.reset() - nodes = [dict(id='0', kind='Code', name='c1', args='def test(x: int): return x'), - dict(id='1', kind='Code', name='c2', args='def test(x: int): return 2 * x'), - dict(id='2', kind='Code', name='c3', args='def test(x: int): return 3 * x'), - dict(id='3', kind='JoinFormatter', name='join', args=dict(type='sum'))] + nodes = [ + dict( + id="0", + kind="Code", + name="c1", + args=dict(code="def test(x: int): return x"), + ), + dict( + id="1", + kind="Code", + name="c2", + args=dict(code="def test(x: int): return 2 * x"), + ), + dict( + id="2", + kind="Code", + name="c3", + args=dict(code="def test(x: int): return 3 * x"), + ), + dict(id="3", kind="JoinFormatter", name="join", args=dict(type="sum")), + ] edges = [dict(iid='__start__', oid='0'), dict(iid='__start__', oid='1'), dict(iid='__start__', oid='2'), dict(iid='0', oid='3'), dict(iid='1', oid='3'), dict(iid='2', oid='3'), dict(iid='3', oid='__end__')] gid = engine.start(nodes, edges) @@ -322,10 +514,32 @@ def test_engine_join_sum(self): assert engine.run(gid, [1]) == [1, 1, 1, 1, 1, 1] def test_engine_join_todict(self): - nodes = [dict(id='0', kind='Code', name='c1', args='def test(x: int): return x'), - dict(id='1', kind='Code', name='c2', args='def test(x: int): return 2 * x'), - dict(id='2', kind='Code', name='c3', args='def test(x: int): return 3 * x'), - dict(id='3', kind='JoinFormatter', name='join', args=dict(type='to_dict', names=['a', 'b', 'c']))] + nodes = [ + dict( + id="0", + kind="Code", + name="c1", + args=dict(code="def test(x: int): return x"), + ), + dict( + id="1", + kind="Code", + name="c2", + args=dict(code="def test(x: int): return 2 * x"), + ), + dict( + id="2", + kind="Code", + name="c3", + args=dict(code="def test(x: int): return 3 * x"), + ), + dict( + id="3", + kind="JoinFormatter", + name="join", + args=dict(type="to_dict", names=["a", "b", "c"]), + ), + ] edges = [dict(iid='__start__', oid='0'), dict(iid='__start__', oid='1'), dict(iid='__start__', oid='2'), dict(iid='0', oid='3'), dict(iid='1', oid='3'), dict(iid='2', oid='3'), dict(iid='3', oid='__end__')] engine = LightEngine() @@ -335,9 +549,24 @@ def test_engine_join_todict(self): assert engine.run(gid, [1]) == dict(a=[1], b=[1, 1], c=[1, 1, 1]) def test_engine_update(self): - plus1 = dict(id='1', kind='Code', name='m1', args='def test(x: int):\n return 1 + x\n') - double = dict(id='2', kind='Code', name='m2', args='def test(x: int):\n return 2 * x\n') - square = dict(id='3', kind='Code', name='m3', args='def test(x: int):\n return x * x\n') + plus1 = dict( + id="1", + kind="Code", + name="m1", + args=dict(code="def test(x: int):\n return 1 + x\n"), + ) + double = dict( + id="2", + kind="Code", + name="m2", + args=dict(code="def test(x: int):\n return 2 * x\n"), + ) + square = dict( + id="3", + kind="Code", + name="m3", + args=dict(code="def test(x: int):\n return x * x\n"), + ) ifs = dict(id='4', kind='Ifs', name='i1', args=dict( cond='def cond(x): return x < 10', true=[plus1, double], false=[square] )) @@ -349,7 +578,12 @@ def test_engine_update(self): assert engine.run(gid, 5) == 12 assert engine.run(gid, 10) == 100 - double = dict(id='2', kind='Code', name='m2', args='def test(x: int):\n return 3 * x\n') + double = dict( + id="2", + kind="Code", + name="m2", + args=dict(code="def test(x: int):\n return 3 * x\n"), + ) ifs = dict(id='4', kind='Ifs', name='i1', args=dict( cond='def cond(x): return x < 10', true=[plus1, double], false=[square] )) @@ -361,10 +595,27 @@ def test_engine_update(self): assert engine.run(gid, 10) == 100 def test_engine_join_join(self): - nodes = [dict(id='0', kind='Code', name='c1', args='def test(x: int): return x'), - dict(id='1', kind='Code', name='c2', args='def test(x: int): return 2 * x'), - dict(id='2', kind='Code', name='c3', args='def test(x: int): return 3 * x'), - dict(id='3', kind='JoinFormatter', name='join', args=dict(type='join'))] + nodes = [ + dict( + id="0", + kind="Code", + name="c1", + args=dict(code="def test(x: int): return x"), + ), + dict( + id="1", + kind="Code", + name="c2", + args=dict(code="def test(x: int): return 2 * x"), + ), + dict( + id="2", + kind="Code", + name="c3", + args=dict(code="def test(x: int): return 3 * x"), + ), + dict(id="3", kind="JoinFormatter", name="join", args=dict(type="join")), + ] edges = [dict(iid='__start__', oid='0'), dict(iid='__start__', oid='1'), dict(iid='__start__', oid='2'), dict(iid='0', oid='3'), dict(iid='1', oid='3'), dict(iid='2', oid='3'), dict(iid='3', oid='__end__')] engine = LightEngine() @@ -376,7 +627,14 @@ def test_engine_join_join(self): assert engine.run(gid, '1') == '1\n11\n111' def test_engine_server(self): - nodes = [dict(id='1', kind='Code', name='m1', args='def test(x: int):\n return 2 * x\n')] + nodes = [ + dict( + id="1", + kind="Code", + name="m1", + args=dict(code="def test(x: int):\n return 2 * x\n"), + ) + ] edges = [dict(iid='__start__', oid='1'), dict(iid='1', oid='__end__')] resources = [dict(id='2', kind='server', name='s1', args=dict(port=None)), dict(id='3', kind='web', name='w1', args=dict(port=None, title='网页', history=[], audio=False)) @@ -433,11 +691,36 @@ def test_engine_httptool(self): url = 'https://postman-echo.com/get' nodes = [ - dict(id='0', kind='Code', name='code1', args='def p1(): return "foo"'), - dict(id='1', kind='Code', name='code2', args='def p2(): return "bar"'), - dict(id='2', kind='Code', name='code3', args='def h1(): return "baz"'), - dict(id='3', kind='HttpTool', name='http', args=dict( - method='GET', url=url, params=params, headers=headers, _lazyllm_arg_names=['p1', 'p2', 'h1'])) + dict( + id="0", + kind="Code", + name="code1", + args=dict(code='def p1(): return "foo"'), + ), + dict( + id="1", + kind="Code", + name="code2", + args=dict(code='def p2(): return "bar"'), + ), + dict( + id="2", + kind="Code", + name="code3", + args=dict(code='def h1(): return "baz"'), + ), + dict( + id="3", + kind="HttpTool", + name="http", + args=dict( + method="GET", + url=url, + params=params, + headers=headers, + _lazyllm_arg_names=["p1", "p2", "h1"], + ), + ), ] edges = [dict(iid='__start__', oid='0'), dict(iid='__start__', oid='1'), dict(iid='__start__', oid='2'), dict(iid='0', oid='3'), dict(iid='1', oid='3'), dict(iid='2', oid='3'), dict(iid='3', oid='__end__')] @@ -454,9 +737,24 @@ def test_engine_status(self): resources = [dict(id='0', kind='LocalLLM', name='m1', args=dict(base_model='', deploy_method='dummy'))] llm_node = dict(id='1', kind='SharedLLM', name='s1', args=dict(llm='0', prompt=None)) - plus1 = dict(id='2', kind='Code', name='m1', args='def test(x: int):\n return 1 + x\n') - double = dict(id='3', kind='Code', name='m2', args='def test(x: int):\n return 2 * x\n') - square = dict(id='4', kind='Code', name='m3', args='def test(x: int):\n return x * x\n') + plus1 = dict( + id="2", + kind="Code", + name="m1", + args=dict(code="def test(x: int):\n return 1 + x\n"), + ) + double = dict( + id="3", + kind="Code", + name="m2", + args=dict(code="def test(x: int):\n return 2 * x\n"), + ) + square = dict( + id="4", + kind="Code", + name="m3", + args=dict(code="def test(x: int):\n return x * x\n"), + ) subgraph = dict(id='5', kind='SubGraph', name='subgraph', args=dict(nodes=[double, plus1])) ifs = dict(id='6', kind='Ifs', name='i1', args=dict( @@ -490,14 +788,41 @@ def test_rag(self): resources = [ dict(id='00', kind='LocalEmbedding', name='e1', args=dict(base_model='bge-large-zh-v1.5')), dict(id='0', kind='Document', name='d1', args=dict(dataset_path='rag_master', embed='00'))] - nodes = [dict(id='1', kind='Retriever', name='ret1', - args=dict(doc='0', group_name='CoarseChunk', similarity='bm25_chinese', topk=3)), - dict(id='4', kind='Reranker', name='rek1', - args=dict(type='ModuleReranker', output_format='content', join=True, - arguments=dict(model="bge-reranker-large", topk=1))), - dict(id='5', kind='Code', name='c1', - args='def test(nodes, query): return f\'context_str={nodes}, query={query}\''), - dict(id='6', kind='LocalLLM', name='m1', args=dict(base_model='', deploy_method='dummy'))] + nodes = [ + dict( + id="1", + kind="Retriever", + name="ret1", + args=dict( + doc="0", group_name="CoarseChunk", similarity="bm25_chinese", topk=3 + ), + ), + dict( + id="4", + kind="Reranker", + name="rek1", + args=dict( + type="ModuleReranker", + output_format="content", + join=True, + arguments=dict(model="bge-reranker-large", topk=1), + ), + ), + dict( + id="5", + kind="Code", + name="c1", + args=dict( + code="def test(nodes, query): return f'context_str={nodes}, query={query}'" + ), + ), + dict( + id="6", + kind="LocalLLM", + name="m1", + args=dict(base_model="", deploy_method="dummy"), + ), + ] edges = [dict(iid='__start__', oid='1'), dict(iid='1', oid='4'), dict(iid='__start__', oid='4'), dict(iid='4', oid='5'), dict(iid='__start__', oid='5'), dict(iid='5', oid='6'), dict(iid='6', oid='__end__')] From 09a059eaac5174ff4c6e672b0558757085f7f473 Mon Sep 17 00:00:00 2001 From: zhangyongchao Date: Tue, 12 Nov 2024 15:45:44 +0800 Subject: [PATCH 02/12] refactor --- tests/basic_tests/test_engine.py | 642 ++++++++++------------------ tests/charge_tests/test_sql_tool.py | 7 - 2 files changed, 229 insertions(+), 420 deletions(-) diff --git a/tests/basic_tests/test_engine.py b/tests/basic_tests/test_engine.py index cfab0269..5946ad45 100644 --- a/tests/basic_tests/test_engine.py +++ b/tests/basic_tests/test_engine.py @@ -60,14 +60,7 @@ def test_engine_subgraph(self): assert '1234' in r def test_engine_code(self): - nodes = [ - dict( - id="1", - kind="Code", - name="m1", - args=dict(code="def test(x: int):\n return 2 * x\n"), - ) - ] + nodes = [dict(id='1', kind='Code', name='m1', args=dict(code='def test(x: int):\n return 2 * x\n'))] edges = [dict(iid='__start__', oid='1'), dict(iid='1', oid='__end__')] engine = LightEngine() @@ -76,35 +69,23 @@ def test_engine_code(self): assert engine.run(gid, 2) == 4 def test_engine_switch(self): - plus1 = dict( - id="1", - kind="Code", - name="m1", - args=dict(code="def test(x: int):\n return 1 + x\n"), - ) - double = dict( - id="2", - kind="Code", - name="m2", - args=dict(code="def test(x: int):\n return 2 * x\n"), - ) - # square = dict(id='3', kind='Code', name='m3', args='def test(x: int):\n return x * x\n') - square = dict( - id="3", - kind="Code", - name="m3", - args=dict( - code="def test(x: int):\n return x * x\n", - _lazyllm_enable_report=True, - ), - ) + plus1 = dict(id='1', kind='Code', name='m1', args=dict(code='def test(x: int):\n return 1 + x\n')) + double = dict(id='2', kind='Code', name='m2', args=dict(code='def test(x: int):\n return 2 * x\n')) + square = dict(id='3', + kind='Code', + name='m3', + args=dict(code='def test(x: int):\n return x * x\n', _lazyllm_enable_report=True)) switch = dict( id="4", kind="Switch", name="s1", args=dict( judge_on_full_input=True, - nodes={1: [double], 2: [plus1, double], 3: [square]}, + nodes={ + 1: [double], + 2: [plus1, double], + 3: [square] + }, _lazyllm_enable_report=True, ), ) @@ -118,7 +99,6 @@ def test_engine_switch(self): assert engine.run(gid, 3) == 9 engine.reset() - lazyllm.globals._init_sid("session_test_002") switch = dict( id="4", @@ -126,7 +106,11 @@ def test_engine_switch(self): name="s1", args=dict( judge_on_full_input=False, - nodes={"case1": [double], "case2": [plus1, double], "case3": [square]}, + nodes={ + "case1": [double], + "case2": [plus1, double], + "case3": [square] + }, _lazyllm_enable_report=True, ), ) @@ -139,28 +123,9 @@ def test_engine_switch(self): assert engine.run(gid, 'case3', 3) == 9 def test_engine_ifs(self): - # plus1 = dict(id='1', kind='Code', name='m1', args='def test(x: int):\n return 1 + x\n') - plus1 = dict( - id="1", - kind="Code", - name="m1", - args=dict( - code="def test(x: int):\n return 1 + x\n", - _lazyllm_enable_report=True, - ), - ) - double = dict( - id="2", - kind="Code", - name="m2", - args=dict(code="def test(x: int):\n return 2 * x\n"), - ) - square = dict( - id="3", - kind="Code", - name="m3", - args=dict(code="def test(x: int):\n return x * x\n"), - ) + plus1 = dict(id='1', kind='Code', name='m1', args=dict(code='def test(x: int):\n return 1 + x\n')) + double = dict(id='2', kind='Code', name='m2', args=dict(code='def test(x: int):\n return 2 * x\n')) + square = dict(id='3', kind='Code', name='m3', args=dict(code='def test(x: int):\n return x * x\n')) ifs = dict( id="4", kind="Ifs", @@ -182,14 +147,7 @@ def test_engine_ifs(self): assert engine.run(gid, 10) == 100 def test_engine_loop(self): - nodes = [ - dict( - id="1", - kind="Code", - name="code", - args=dict(code="def square(x: int): return x * x"), - ) - ] + nodes = [dict(id='1', kind='Code', name='code', args=dict(code='def square(x: int): return x * x'))] edges = [dict(iid='__start__', oid='1'), dict(iid='1', oid='__end__')] nodes = [ @@ -213,14 +171,7 @@ def test_engine_loop(self): assert engine.run(gid, 2) == 16 def test_engine_warp(self): - nodes = [ - dict( - id="1", - kind="Code", - name="code", - args=dict(code="def square(x: int): return x * x"), - ) - ] + nodes = [dict(id='1', kind='Code', name='code', args=dict(code='def square(x: int): return x * x'))] edges = [dict(iid='__start__', oid='1'), dict(iid='1', oid='__end__')] nodes = [ @@ -283,36 +234,23 @@ def test_engine_formatter(self): def test_engine_edge_formatter(self): nodes = [ - dict( - id="1", - kind="Code", - name="m1", - args=dict(code="def test(x: int):\n return x\n"), - ), - dict( - id="2", - kind="Code", - name="m2", - args=dict( - code="def test(x: int):\n return [[x, 2*x], [3*x, 4*x]]\n" - ), - ), - dict( - id="3", - kind="Code", - name="m3", - args=dict(code="def test(x: int):\n return dict(a=1, b=x * x)\n"), - ), - dict( - id="4", - kind="Code", - name="m4", - args=dict(code='def test(x, y, z):\n return f"{x}{y}{z}"\n'), - ), + dict(id='1', kind='Code', name='m1', args=dict(code='def test(x: int):\n return x\n')), + dict(id='2', + kind='Code', + name='m2', + args=dict(code='def test(x: int):\n return [[x, 2*x], [3*x, 4*x]]\n')), + dict(id='3', kind='Code', name='m3', args=dict(code='def test(x: int):\n return dict(a=1, b=x * x)\n')), + dict(id='4', kind='Code', name='m4', args=dict(code='def test(x, y, z):\n return f"{x}{y}{z}"\n')) + ] + edges = [ + dict(iid='__start__', oid='1'), + dict(iid='__start__', oid='2'), + dict(iid='__start__', oid='3'), + dict(iid='1', oid='4'), + dict(iid='2', oid='4', formatter='[:][1]'), + dict(iid='3', oid='4', formatter='[b]'), + dict(iid='4', oid='__end__') ] - edges = [dict(iid='__start__', oid='1'), dict(iid='__start__', oid='2'), dict(iid='__start__', oid='3'), - dict(iid='1', oid='4'), dict(iid='2', oid='4', formatter='[:][1]'), - dict(iid='3', oid='4', formatter='[b]'), dict(iid='4', oid='__end__')] engine = LightEngine() gid = engine.start(nodes, edges) @@ -321,27 +259,17 @@ def test_engine_edge_formatter(self): def test_engine_edge_formatter_start(self): nodes = [ - dict( - id="1", - kind="Code", - name="m1", - args=dict(code="def test(x: int): return x"), - ), - dict( - id="2", - kind="Code", - name="m2", - args=dict(code="def test(x: int): return 2 * x"), - ), - dict( - id="3", - kind="Code", - name="m3", - args=dict(code="def test(x, y): return x + y"), - ), + dict(id='1', kind='Code', name='m1', args=dict(code='def test(x: int): return x')), + dict(id='2', kind='Code', name='m2', args=dict(code='def test(x: int): return 2 * x')), + dict(id='3', kind='Code', name='m3', args=dict(code='def test(x, y): return x + y')) + ] + edges = [ + dict(iid='__start__', oid='1', formatter='[0]'), + dict(iid='__start__', oid='2', formatter='[1]'), + dict(iid='1', oid='3'), + dict(iid='2', oid='3'), + dict(iid='3', oid='__end__') ] - edges = [dict(iid='__start__', oid='1', formatter='[0]'), dict(iid='__start__', oid='2', formatter='[1]'), - dict(iid='1', oid='3'), dict(iid='2', oid='3'), dict(iid='3', oid='__end__')] engine = LightEngine() gid = engine.start(nodes, edges) @@ -350,36 +278,22 @@ def test_engine_edge_formatter_start(self): def test_engine_formatter_end(self): nodes = [ - dict( - id="1", - kind="Code", - name="m1", - args=dict(code="def test(x: int):\n return x\n"), - ), - dict( - id="2", - kind="Code", - name="m2", - args=dict( - code="def test1(x: int):\n return [[x, 2*x], [3*x, 4*x]]\n" - ), - ), + dict(id='1', kind='Code', name='m1', args=dict(code='def test(x: int):\n return x\n')), + dict(id='2', + kind='Code', + name='m2', + args=dict(code='def test1(x: int):\n return [[x, 2*x], [3*x, 4*x]]\n')), # two unused node - dict( - id="3", - kind="Code", - name="m3", - args=dict(code="def test2(x: int):\n return dict(a=1, b=x * x)\n"), - ), - dict( - id="4", - kind="Code", - name="m4", - args=dict(code='def test3(x, y, z):\n return f"{x}{y}{z}"\n'), - ), + dict(id='3', kind='Code', name='m3', + args=dict(code='def test2(x: int):\n return dict(a=1, b=x * x)\n')), + dict(id='4', kind='Code', name='m4', args=dict(code='def test3(x, y, z):\n return f"{x}{y}{z}"\n')) + ] + edges = [ + dict(iid='__start__', oid='1'), + dict(iid='__start__', oid='2'), + dict(iid='2', oid='__end__'), + dict(iid='1', oid='__end__') ] - edges = [dict(iid='__start__', oid='1'), dict(iid='__start__', oid='2'), dict(iid='2', oid='__end__'), - dict(iid='1', oid='__end__')] engine = LightEngine() gid = engine.start(nodes, edges) @@ -390,29 +304,20 @@ def test_engine_formatter_end(self): engine.reset() nodes = [ - dict( - id="1", - kind="Code", - name="m1", - args=dict(code="def test(x: int):\n return x\n"), - ), - dict( - id="2", - kind="Code", - name="m2", - args=dict( - code="def test1(x: int):\n return [[x, 2*x], [3*x, 4*x]]\n" - ), - ), - dict( - id="3", - kind="JoinFormatter", - name="join", - args=dict(type="to_dict", names=["a", "b"]), - ), + dict(id='1', kind='Code', name='m1', args=dict(code='def test(x: int):\n return x\n')), + dict(id='2', + kind='Code', + name='m2', + args=dict(code='def test1(x: int):\n return [[x, 2*x], [3*x, 4*x]]\n')), + dict(id='3', kind='JoinFormatter', name='join', args=dict(type='to_dict', names=['a', 'b'])) + ] + edges = [ + dict(iid='__start__', oid='1'), + dict(iid='__start__', oid='2'), + dict(iid='2', oid='3'), + dict(iid='1', oid='3'), + dict(iid='3', oid='__end__', formatter='*[a, b]') ] - edges = [dict(iid='__start__', oid='1'), dict(iid='__start__', oid='2'), dict(iid='2', oid='3'), - dict(iid='1', oid='3'), dict(iid='3', oid='__end__', formatter='*[a, b]')] engine = LightEngine() gid = engine.start(nodes, edges) r = engine.run(gid, 1) @@ -421,13 +326,8 @@ def test_engine_formatter_end(self): def test_engine_join_stack(self): nodes = [ - dict( - id="0", - kind="Code", - name="c1", - args=dict(code="def test(x: int): return x"), - ), - dict(id="1", kind="JoinFormatter", name="join", args=dict(type="stack")), + dict(id='0', kind='Code', name='c1', args=dict(code='def test(x: int): return x')), + dict(id='1', kind='JoinFormatter', name='join', args=dict(type='stack')) ] edges = [dict(iid='__start__', oid='0'), dict(iid='0', oid='1'), dict(iid='1', oid='__end__')] engine = LightEngine() @@ -439,28 +339,20 @@ def test_engine_join_stack(self): engine.reset() nodes = [ - dict( - id="0", - kind="Code", - name="c1", - args=dict(code="def test(x: int): return x"), - ), - dict( - id="1", - kind="Code", - name="c2", - args=dict(code="def test(x: int): return 2 * x"), - ), - dict( - id="2", - kind="Code", - name="c3", - args=dict(code="def test(x: int): return 3 * x"), - ), - dict(id="3", kind="JoinFormatter", name="join", args=dict(type="stack")), + dict(id='0', kind='Code', name='c1', args=dict(code='def test(x: int): return x')), + dict(id='1', kind='Code', name='c2', args=dict(code='def test(x: int): return 2 * x')), + dict(id='2', kind='Code', name='c3', args=dict(code='def test(x: int): return 3 * x')), + dict(id='3', kind='JoinFormatter', name='join', args=dict(type='stack')) + ] + edges = [ + dict(iid='__start__', oid='0'), + dict(iid='__start__', oid='1'), + dict(iid='__start__', oid='2'), + dict(iid='0', oid='3'), + dict(iid='1', oid='3'), + dict(iid='2', oid='3'), + dict(iid='3', oid='__end__') ] - edges = [dict(iid='__start__', oid='0'), dict(iid='__start__', oid='1'), dict(iid='__start__', oid='2'), - dict(iid='0', oid='3'), dict(iid='1', oid='3'), dict(iid='2', oid='3'), dict(iid='3', oid='__end__')] gid = engine.start(nodes, edges) assert engine.run(gid, 1) == [1, 2, 3] assert engine.run(gid, '1') == ['1', '11', '111'] @@ -468,13 +360,8 @@ def test_engine_join_stack(self): def test_engine_join_sum(self): nodes = [ - dict( - id="0", - kind="Code", - name="c1", - args=dict(code="def test(x: int): return [x, 2 * x]"), - ), - dict(id="1", kind="JoinFormatter", name="join", args=dict(type="sum")), + dict(id='0', kind='Code', name='c1', args=dict(code='def test(x: int): return [x, 2 * x]')), + dict(id='1', kind='JoinFormatter', name='join', args=dict(type='sum')) ] edges = [dict(iid='__start__', oid='0'), dict(iid='0', oid='1'), dict(iid='1', oid='__end__')] engine = LightEngine() @@ -486,28 +373,20 @@ def test_engine_join_sum(self): engine.reset() nodes = [ - dict( - id="0", - kind="Code", - name="c1", - args=dict(code="def test(x: int): return x"), - ), - dict( - id="1", - kind="Code", - name="c2", - args=dict(code="def test(x: int): return 2 * x"), - ), - dict( - id="2", - kind="Code", - name="c3", - args=dict(code="def test(x: int): return 3 * x"), - ), - dict(id="3", kind="JoinFormatter", name="join", args=dict(type="sum")), + dict(id='0', kind='Code', name='c1', args=dict(code='def test(x: int): return x')), + dict(id='1', kind='Code', name='c2', args=dict(code='def test(x: int): return 2 * x')), + dict(id='2', kind='Code', name='c3', args=dict(code='def test(x: int): return 3 * x')), + dict(id='3', kind='JoinFormatter', name='join', args=dict(type='sum')) + ] + edges = [ + dict(iid='__start__', oid='0'), + dict(iid='__start__', oid='1'), + dict(iid='__start__', oid='2'), + dict(iid='0', oid='3'), + dict(iid='1', oid='3'), + dict(iid='2', oid='3'), + dict(iid='3', oid='__end__') ] - edges = [dict(iid='__start__', oid='0'), dict(iid='__start__', oid='1'), dict(iid='__start__', oid='2'), - dict(iid='0', oid='3'), dict(iid='1', oid='3'), dict(iid='2', oid='3'), dict(iid='3', oid='__end__')] gid = engine.start(nodes, edges) assert engine.run(gid, 1) == 6 assert engine.run(gid, '1') == '111111' @@ -515,33 +394,20 @@ def test_engine_join_sum(self): def test_engine_join_todict(self): nodes = [ - dict( - id="0", - kind="Code", - name="c1", - args=dict(code="def test(x: int): return x"), - ), - dict( - id="1", - kind="Code", - name="c2", - args=dict(code="def test(x: int): return 2 * x"), - ), - dict( - id="2", - kind="Code", - name="c3", - args=dict(code="def test(x: int): return 3 * x"), - ), - dict( - id="3", - kind="JoinFormatter", - name="join", - args=dict(type="to_dict", names=["a", "b", "c"]), - ), + dict(id='0', kind='Code', name='c1', args=dict(code='def test(x: int): return x')), + dict(id='1', kind='Code', name='c2', args=dict(code='def test(x: int): return 2 * x')), + dict(id='2', kind='Code', name='c3', args=dict(code='def test(x: int): return 3 * x')), + dict(id='3', kind='JoinFormatter', name='join', args=dict(type='to_dict', names=['a', 'b', 'c'])) + ] + edges = [ + dict(iid='__start__', oid='0'), + dict(iid='__start__', oid='1'), + dict(iid='__start__', oid='2'), + dict(iid='0', oid='3'), + dict(iid='1', oid='3'), + dict(iid='2', oid='3'), + dict(iid='3', oid='__end__') ] - edges = [dict(iid='__start__', oid='0'), dict(iid='__start__', oid='1'), dict(iid='__start__', oid='2'), - dict(iid='0', oid='3'), dict(iid='1', oid='3'), dict(iid='2', oid='3'), dict(iid='3', oid='__end__')] engine = LightEngine() gid = engine.start(nodes, edges) assert engine.run(gid, 1) == dict(a=1, b=2, c=3) @@ -549,27 +415,13 @@ def test_engine_join_todict(self): assert engine.run(gid, [1]) == dict(a=[1], b=[1, 1], c=[1, 1, 1]) def test_engine_update(self): - plus1 = dict( - id="1", - kind="Code", - name="m1", - args=dict(code="def test(x: int):\n return 1 + x\n"), - ) - double = dict( - id="2", - kind="Code", - name="m2", - args=dict(code="def test(x: int):\n return 2 * x\n"), - ) - square = dict( - id="3", - kind="Code", - name="m3", - args=dict(code="def test(x: int):\n return x * x\n"), - ) - ifs = dict(id='4', kind='Ifs', name='i1', args=dict( - cond='def cond(x): return x < 10', true=[plus1, double], false=[square] - )) + plus1 = dict(id='1', kind='Code', name='m1', args=dict(code='def test(x: int):\n return 1 + x\n')) + double = dict(id='2', kind='Code', name='m2', args=dict(code='def test(x: int):\n return 2 * x\n')) + square = dict(id='3', kind='Code', name='m3', args=dict(code='def test(x: int):\n return x * x\n')) + ifs = dict(id='4', + kind='Ifs', + name='i1', + args=dict(cond='def cond(x): return x < 10', true=[plus1, double], false=[square])) nodes = [ifs] edges = [dict(iid='__start__', oid='4'), dict(iid='4', oid='__end__')] engine = LightEngine() @@ -578,15 +430,11 @@ def test_engine_update(self): assert engine.run(gid, 5) == 12 assert engine.run(gid, 10) == 100 - double = dict( - id="2", - kind="Code", - name="m2", - args=dict(code="def test(x: int):\n return 3 * x\n"), - ) - ifs = dict(id='4', kind='Ifs', name='i1', args=dict( - cond='def cond(x): return x < 10', true=[plus1, double], false=[square] - )) + double = dict(id='2', kind='Code', name='m2', args=dict(code='def test(x: int):\n return 3 * x\n')) + ifs = dict(id='4', + kind='Ifs', + name='i1', + args=dict(cond='def cond(x): return x < 10', true=[plus1, double], false=[square])) nodes = [ifs] engine.update(gid, nodes, edges) @@ -596,28 +444,20 @@ def test_engine_update(self): def test_engine_join_join(self): nodes = [ - dict( - id="0", - kind="Code", - name="c1", - args=dict(code="def test(x: int): return x"), - ), - dict( - id="1", - kind="Code", - name="c2", - args=dict(code="def test(x: int): return 2 * x"), - ), - dict( - id="2", - kind="Code", - name="c3", - args=dict(code="def test(x: int): return 3 * x"), - ), - dict(id="3", kind="JoinFormatter", name="join", args=dict(type="join")), + dict(id='0', kind='Code', name='c1', args=dict(code='def test(x: int): return x')), + dict(id='1', kind='Code', name='c2', args=dict(code='def test(x: int): return 2 * x')), + dict(id='2', kind='Code', name='c3', args=dict(code='def test(x: int): return 3 * x')), + dict(id='3', kind='JoinFormatter', name='join', args=dict(type='join')) + ] + edges = [ + dict(iid='__start__', oid='0'), + dict(iid='__start__', oid='1'), + dict(iid='__start__', oid='2'), + dict(iid='0', oid='3'), + dict(iid='1', oid='3'), + dict(iid='2', oid='3'), + dict(iid='3', oid='__end__') ] - edges = [dict(iid='__start__', oid='0'), dict(iid='__start__', oid='1'), dict(iid='__start__', oid='2'), - dict(iid='0', oid='3'), dict(iid='1', oid='3'), dict(iid='2', oid='3'), dict(iid='3', oid='__end__')] engine = LightEngine() gid = engine.start(nodes, edges) assert engine.run(gid, '1') == '111111' @@ -627,18 +467,12 @@ def test_engine_join_join(self): assert engine.run(gid, '1') == '1\n11\n111' def test_engine_server(self): - nodes = [ - dict( - id="1", - kind="Code", - name="m1", - args=dict(code="def test(x: int):\n return 2 * x\n"), - ) - ] + nodes = [dict(id='1', kind='Code', name='m1', args=dict(code='def test(x: int):\n return 2 * x\n'))] edges = [dict(iid='__start__', oid='1'), dict(iid='1', oid='__end__')] - resources = [dict(id='2', kind='server', name='s1', args=dict(port=None)), - dict(id='3', kind='web', name='w1', args=dict(port=None, title='网页', history=[], audio=False)) - ] + resources = [ + dict(id='2', kind='server', name='s1', args=dict(port=None)), + dict(id='3', kind='web', name='w1', args=dict(port=None, title='网页', history=[], audio=False)) + ] engine = LightEngine() gid = engine.start(nodes, edges, resources, gid='graph-1') assert engine.status(gid) == {'1': 'running', '2': lazyllm.launcher.Status.Running, '3': 'running'} @@ -691,39 +525,27 @@ def test_engine_httptool(self): url = 'https://postman-echo.com/get' nodes = [ - dict( - id="0", - kind="Code", - name="code1", - args=dict(code='def p1(): return "foo"'), - ), - dict( - id="1", - kind="Code", - name="code2", - args=dict(code='def p2(): return "bar"'), - ), - dict( - id="2", - kind="Code", - name="code3", - args=dict(code='def h1(): return "baz"'), - ), - dict( - id="3", - kind="HttpTool", - name="http", - args=dict( - method="GET", - url=url, - params=params, - headers=headers, - _lazyllm_arg_names=["p1", "p2", "h1"], - ), - ), + dict(id='0', kind='Code', name='code1', args=dict(code='def p1(): return "foo"')), + dict(id='1', kind='Code', name='code2', args=dict(code='def p2(): return "bar"')), + dict(id='2', kind='Code', name='code3', args=dict(code='def h1(): return "baz"')), + dict(id='3', + kind='HttpTool', + name='http', + args=dict(method='GET', + url=url, + params=params, + headers=headers, + _lazyllm_arg_names=['p1', 'p2', 'h1'])) + ] + edges = [ + dict(iid='__start__', oid='0'), + dict(iid='__start__', oid='1'), + dict(iid='__start__', oid='2'), + dict(iid='0', oid='3'), + dict(iid='1', oid='3'), + dict(iid='2', oid='3'), + dict(iid='3', oid='__end__') ] - edges = [dict(iid='__start__', oid='0'), dict(iid='__start__', oid='1'), dict(iid='__start__', oid='2'), - dict(iid='0', oid='3'), dict(iid='1', oid='3'), dict(iid='2', oid='3'), dict(iid='3', oid='__end__')] engine = LightEngine() gid = engine.start(nodes, edges, gid='graph-1') @@ -737,33 +559,29 @@ def test_engine_status(self): resources = [dict(id='0', kind='LocalLLM', name='m1', args=dict(base_model='', deploy_method='dummy'))] llm_node = dict(id='1', kind='SharedLLM', name='s1', args=dict(llm='0', prompt=None)) - plus1 = dict( - id="2", - kind="Code", - name="m1", - args=dict(code="def test(x: int):\n return 1 + x\n"), - ) - double = dict( - id="3", - kind="Code", - name="m2", - args=dict(code="def test(x: int):\n return 2 * x\n"), - ) - square = dict( - id="4", - kind="Code", - name="m3", - args=dict(code="def test(x: int):\n return x * x\n"), - ) + plus1 = dict(id='2', kind='Code', name='m1', args=dict(code='def test(x: int):\n return 1 + x\n')) + double = dict(id='3', kind='Code', name='m2', args=dict(code='def test(x: int):\n return 2 * x\n')) + square = dict(id='4', kind='Code', name='m3', args=dict(code='def test(x: int):\n return x * x\n')) subgraph = dict(id='5', kind='SubGraph', name='subgraph', args=dict(nodes=[double, plus1])) - ifs = dict(id='6', kind='Ifs', name='i1', args=dict( - cond='def cond(x): return x % 2 == 0', true=plus1, false=[square])) - loop = dict(id='7', kind='Loop', name='loop', args=dict( - stop_condition='def cond(x): return x > 8', nodes=[double])) - - switch = dict(id='8', kind='Switch', name='sw1', args=dict(judge_on_full_input=True, nodes={ - 1: [plus1, subgraph], 2: ifs, 3: loop, 5: [ifs]})) + ifs = dict(id='6', + kind='Ifs', + name='i1', + args=dict(cond='def cond(x): return x % 2 == 0', true=plus1, false=[square])) + loop = dict(id='7', + kind='Loop', + name='loop', + args=dict(stop_condition='def cond(x): return x > 8', nodes=[double])) + + switch = dict(id='8', + kind='Switch', + name='sw1', + args=dict(judge_on_full_input=True, nodes={ + 1: [plus1, subgraph], + 2: ifs, + 3: loop, + 5: [ifs] + })) warp = dict(id='9', kind='Warp', name='w1', args=dict(nodes=[switch, plus1])) join = dict(id='10', kind='JoinFormatter', name='join', args=dict(type='join', symbol=', ')) @@ -772,14 +590,28 @@ def test_engine_status(self): gid = engine.start(nodes, [], resources) assert '6, 4, 13, 26' in engine.run(gid, 1, 2, 3, 5) - assert engine.status(gid) == {'9': {'8': {'2': 'running', - '5': {'3': 'running', '2': 'running'}, - '6': {'2': 'running', '4': 'running'}, - '7': {'3': 'running'}}, - '2': 'running'}, - '10': 'running', - '1': 'running', - '0': lazyllm.launcher.Status.Running} + assert engine.status(gid) == { + '9': { + '8': { + '2': 'running', + '5': { + '3': 'running', + '2': 'running' + }, + '6': { + '2': 'running', + '4': 'running' + }, + '7': { + '3': 'running' + } + }, + '2': 'running' + }, + '10': 'running', + '1': 'running', + '0': lazyllm.launcher.Status.Running + } class TestEngineRAG(object): @@ -787,41 +619,25 @@ class TestEngineRAG(object): def test_rag(self): resources = [ dict(id='00', kind='LocalEmbedding', name='e1', args=dict(base_model='bge-large-zh-v1.5')), - dict(id='0', kind='Document', name='d1', args=dict(dataset_path='rag_master', embed='00'))] + dict(id='0', kind='Document', name='d1', args=dict(dataset_path='rag_master', embed='00')) + ] nodes = [ - dict( - id="1", - kind="Retriever", - name="ret1", - args=dict( - doc="0", group_name="CoarseChunk", similarity="bm25_chinese", topk=3 - ), - ), - dict( - id="4", - kind="Reranker", - name="rek1", - args=dict( - type="ModuleReranker", - output_format="content", - join=True, - arguments=dict(model="bge-reranker-large", topk=1), - ), - ), - dict( - id="5", - kind="Code", - name="c1", - args=dict( - code="def test(nodes, query): return f'context_str={nodes}, query={query}'" - ), - ), - dict( - id="6", - kind="LocalLLM", - name="m1", - args=dict(base_model="", deploy_method="dummy"), - ), + dict(id='1', + kind='Retriever', + name='ret1', + args=dict(doc='0', group_name='CoarseChunk', similarity='bm25_chinese', topk=3)), + dict(id='4', + kind='Reranker', + name='rek1', + args=dict(type='ModuleReranker', + output_format='content', + join=True, + arguments=dict(model="bge-reranker-large", topk=1))), + dict(id='5', + kind='Code', + name='c1', + args=dict(code='def test(nodes, query): return f\'context_str={nodes}, query={query}\'')), + dict(id='6', kind='LocalLLM', name='m1', args=dict(base_model='', deploy_method='dummy')) ] edges = [dict(iid='__start__', oid='1'), dict(iid='1', oid='4'), dict(iid='__start__', oid='4'), dict(iid='4', oid='5'), dict(iid='__start__', oid='5'), dict(iid='5', oid='6'), diff --git a/tests/charge_tests/test_sql_tool.py b/tests/charge_tests/test_sql_tool.py index 79e22e79..799cefd0 100644 --- a/tests/charge_tests/test_sql_tool.py +++ b/tests/charge_tests/test_sql_tool.py @@ -116,13 +116,6 @@ def test_llm_query_online(self): str_results = sql_call("去年一整年销售额最多的员工是谁，销售额是多少？") self.assertIn("张三", str_results) - @unittest.skip("Charge test has no scc support") - def test_llm_query_local(self): - local_llm = lazyllm.TrainableModule("qwen2-7b-instruct").deploy_method(lazyllm.deploy.vllm).start() - sql_call = SqlCall(local_llm, self.sql_managers[0], use_llm_for_sql_result=True, return_trace=True) - str_results = sql_call("员工编号是3的人来自哪个部门？") - self.assertIn("销售三部", str_results) - if __name__ == "__main__": unittest.main() From 18e443c032aaaddec55c77d60667579ad623c960 Mon Sep 17 00:00:00 2001 From: zhangyongchao Date: Tue, 12 Nov 2024 17:11:25 +0800 Subject: [PATCH 03/12] refactor --- tests/basic_tests/test_engine.py | 330 +++++++++---------------------- 1 file changed, 97 insertions(+), 233 deletions(-) diff --git a/tests/basic_tests/test_engine.py b/tests/basic_tests/test_engine.py index 5946ad45..167e9ef2 100644 --- a/tests/basic_tests/test_engine.py +++ b/tests/basic_tests/test_engine.py @@ -71,21 +71,14 @@ def test_engine_code(self): def test_engine_switch(self): plus1 = dict(id='1', kind='Code', name='m1', args=dict(code='def test(x: int):\n return 1 + x\n')) double = dict(id='2', kind='Code', name='m2', args=dict(code='def test(x: int):\n return 2 * x\n')) - square = dict(id='3', - kind='Code', - name='m3', - args=dict(code='def test(x: int):\n return x * x\n', _lazyllm_enable_report=True)) + square = dict(id='3', kind='Code', name='m3', args=dict(code='def test(x: int):\n return x * x\n')) switch = dict( id="4", kind="Switch", name="s1", args=dict( judge_on_full_input=True, - nodes={ - 1: [double], - 2: [plus1, double], - 3: [square] - }, + nodes={1: [double], 2: [plus1, double], 3: [square]}, _lazyllm_enable_report=True, ), ) @@ -106,11 +99,7 @@ def test_engine_switch(self): name="s1", args=dict( judge_on_full_input=False, - nodes={ - "case1": [double], - "case2": [plus1, double], - "case3": [square] - }, + nodes={"case1": [double], "case2": [plus1, double], "case3": [square]}, _lazyllm_enable_report=True, ), ) @@ -233,24 +222,15 @@ def test_engine_formatter(self): '{"query": "hiaha", "files": ["path/to/file"]}' def test_engine_edge_formatter(self): - nodes = [ - dict(id='1', kind='Code', name='m1', args=dict(code='def test(x: int):\n return x\n')), - dict(id='2', - kind='Code', - name='m2', - args=dict(code='def test(x: int):\n return [[x, 2*x], [3*x, 4*x]]\n')), - dict(id='3', kind='Code', name='m3', args=dict(code='def test(x: int):\n return dict(a=1, b=x * x)\n')), - dict(id='4', kind='Code', name='m4', args=dict(code='def test(x, y, z):\n return f"{x}{y}{z}"\n')) - ] - edges = [ - dict(iid='__start__', oid='1'), - dict(iid='__start__', oid='2'), - dict(iid='__start__', oid='3'), - dict(iid='1', oid='4'), - dict(iid='2', oid='4', formatter='[:][1]'), - dict(iid='3', oid='4', formatter='[b]'), - dict(iid='4', oid='__end__') - ] + nodes = [dict(id='1', kind='Code', name='m1', args=dict(code='def test(x: int):\n return x\n')), + dict(id='2', kind='Code', name='m2', + args=dict(code='def test(x: int):\n return [[x, 2*x], [3*x, 4*x]]\n')), + dict(id='3', kind='Code', name='m3', + args=dict(code='def test(x: int):\n return dict(a=1, b=x * x)\n')), + dict(id='4', kind='Code', name='m4', args=dict(code='def test(x, y, z):\n return f"{x}{y}{z}"\n'))] + edges = [dict(iid='__start__', oid='1'), dict(iid='__start__', oid='2'), dict(iid='__start__', oid='3'), + dict(iid='1', oid='4'), dict(iid='2', oid='4', formatter='[:][1]'), + dict(iid='3', oid='4', formatter='[b]'), dict(iid='4', oid='__end__')] engine = LightEngine() gid = engine.start(nodes, edges) @@ -258,18 +238,11 @@ def test_engine_edge_formatter(self): assert engine.run(gid, 2) == '2[4, 8]4' def test_engine_edge_formatter_start(self): - nodes = [ - dict(id='1', kind='Code', name='m1', args=dict(code='def test(x: int): return x')), - dict(id='2', kind='Code', name='m2', args=dict(code='def test(x: int): return 2 * x')), - dict(id='3', kind='Code', name='m3', args=dict(code='def test(x, y): return x + y')) - ] - edges = [ - dict(iid='__start__', oid='1', formatter='[0]'), - dict(iid='__start__', oid='2', formatter='[1]'), - dict(iid='1', oid='3'), - dict(iid='2', oid='3'), - dict(iid='3', oid='__end__') - ] + nodes = [dict(id='1', kind='Code', name='m1', args=dict(code='def test(x: int): return x')), + dict(id='2', kind='Code', name='m2', args=dict(code='def test(x: int): return 2 * x')), + dict(id='3', kind='Code', name='m3', args=dict(code='def test(x, y): return x + y'))] + edges = [dict(iid='__start__', oid='1', formatter='[0]'), dict(iid='__start__', oid='2', formatter='[1]'), + dict(iid='1', oid='3'), dict(iid='2', oid='3'), dict(iid='3', oid='__end__')] engine = LightEngine() gid = engine.start(nodes, edges) @@ -277,23 +250,14 @@ def test_engine_edge_formatter_start(self): assert engine.run(gid, 5, 3, 1) == 11 def test_engine_formatter_end(self): - nodes = [ - dict(id='1', kind='Code', name='m1', args=dict(code='def test(x: int):\n return x\n')), - dict(id='2', - kind='Code', - name='m2', - args=dict(code='def test1(x: int):\n return [[x, 2*x], [3*x, 4*x]]\n')), - # two unused node - dict(id='3', kind='Code', name='m3', - args=dict(code='def test2(x: int):\n return dict(a=1, b=x * x)\n')), - dict(id='4', kind='Code', name='m4', args=dict(code='def test3(x, y, z):\n return f"{x}{y}{z}"\n')) - ] - edges = [ - dict(iid='__start__', oid='1'), - dict(iid='__start__', oid='2'), - dict(iid='2', oid='__end__'), - dict(iid='1', oid='__end__') - ] + nodes = [dict(id='1', kind='Code', name='m1', args=dict(code='def test(x: int):\n return x\n')), + dict(id='2', kind='Code', name='m2', + args=dict(code='def test1(x: int):\n return [[x, 2*x], [3*x, 4*x]]\n')), + # two unused node + dict(id='3', kind='Code', name='m3', args=dict(code='def test2(x: int):\n return dict(a=1, b=x * x)\n')), + dict(id='4', kind='Code', name='m4', args=dict(code='def test3(x, y, z):\n return f"{x}{y}{z}"\n'))] + edges = [dict(iid='__start__', oid='1'), dict(iid='__start__', oid='2'), dict(iid='2', oid='__end__'), + dict(iid='1', oid='__end__')] engine = LightEngine() gid = engine.start(nodes, edges) @@ -303,21 +267,12 @@ def test_engine_formatter_end(self): engine.reset() - nodes = [ - dict(id='1', kind='Code', name='m1', args=dict(code='def test(x: int):\n return x\n')), - dict(id='2', - kind='Code', - name='m2', - args=dict(code='def test1(x: int):\n return [[x, 2*x], [3*x, 4*x]]\n')), - dict(id='3', kind='JoinFormatter', name='join', args=dict(type='to_dict', names=['a', 'b'])) - ] - edges = [ - dict(iid='__start__', oid='1'), - dict(iid='__start__', oid='2'), - dict(iid='2', oid='3'), - dict(iid='1', oid='3'), - dict(iid='3', oid='__end__', formatter='*[a, b]') - ] + nodes = [dict(id='1', kind='Code', name='m1', args=dict(code='def test(x: int):\n return x\n')), + dict(id='2', kind='Code', name='m2', + args=dict(code='def test1(x: int):\n return [[x, 2*x], [3*x, 4*x]]\n')), + dict(id='3', kind='JoinFormatter', name='join', args=dict(type='to_dict', names=['a', 'b']))] + edges = [dict(iid='__start__', oid='1'), dict(iid='__start__', oid='2'), dict(iid='2', oid='3'), + dict(iid='1', oid='3'), dict(iid='3', oid='__end__', formatter='*[a, b]')] engine = LightEngine() gid = engine.start(nodes, edges) r = engine.run(gid, 1) @@ -325,10 +280,8 @@ def test_engine_formatter_end(self): print(isinstance(r, lazyllm.package)) def test_engine_join_stack(self): - nodes = [ - dict(id='0', kind='Code', name='c1', args=dict(code='def test(x: int): return x')), - dict(id='1', kind='JoinFormatter', name='join', args=dict(type='stack')) - ] + nodes = [dict(id='0', kind='Code', name='c1', args=dict(code='def test(x: int): return x')), + dict(id='1', kind='JoinFormatter', name='join', args=dict(type='stack'))] edges = [dict(iid='__start__', oid='0'), dict(iid='0', oid='1'), dict(iid='1', oid='__end__')] engine = LightEngine() gid = engine.start(nodes, edges) @@ -338,31 +291,20 @@ def test_engine_join_stack(self): engine.reset() - nodes = [ - dict(id='0', kind='Code', name='c1', args=dict(code='def test(x: int): return x')), - dict(id='1', kind='Code', name='c2', args=dict(code='def test(x: int): return 2 * x')), - dict(id='2', kind='Code', name='c3', args=dict(code='def test(x: int): return 3 * x')), - dict(id='3', kind='JoinFormatter', name='join', args=dict(type='stack')) - ] - edges = [ - dict(iid='__start__', oid='0'), - dict(iid='__start__', oid='1'), - dict(iid='__start__', oid='2'), - dict(iid='0', oid='3'), - dict(iid='1', oid='3'), - dict(iid='2', oid='3'), - dict(iid='3', oid='__end__') - ] + nodes = [dict(id='0', kind='Code', name='c1', args=dict(code='def test(x: int): return x')), + dict(id='1', kind='Code', name='c2', args=dict(code='def test(x: int): return 2 * x')), + dict(id='2', kind='Code', name='c3', args=dict(code='def test(x: int): return 3 * x')), + dict(id='3', kind='JoinFormatter', name='join', args=dict(type='stack'))] + edges = [dict(iid='__start__', oid='0'), dict(iid='__start__', oid='1'), dict(iid='__start__', oid='2'), + dict(iid='0', oid='3'), dict(iid='1', oid='3'), dict(iid='2', oid='3'), dict(iid='3', oid='__end__')] gid = engine.start(nodes, edges) assert engine.run(gid, 1) == [1, 2, 3] assert engine.run(gid, '1') == ['1', '11', '111'] assert engine.run(gid, [1]) == [[1], [1, 1], [1, 1, 1]] def test_engine_join_sum(self): - nodes = [ - dict(id='0', kind='Code', name='c1', args=dict(code='def test(x: int): return [x, 2 * x]')), - dict(id='1', kind='JoinFormatter', name='join', args=dict(type='sum')) - ] + nodes = [dict(id='0', kind='Code', name='c1', args=dict(code='def test(x: int): return [x, 2 * x]')), + dict(id='1', kind='JoinFormatter', name='join', args=dict(type='sum'))] edges = [dict(iid='__start__', oid='0'), dict(iid='0', oid='1'), dict(iid='1', oid='__end__')] engine = LightEngine() gid = engine.start(nodes, edges) @@ -372,42 +314,24 @@ def test_engine_join_sum(self): engine.reset() - nodes = [ - dict(id='0', kind='Code', name='c1', args=dict(code='def test(x: int): return x')), - dict(id='1', kind='Code', name='c2', args=dict(code='def test(x: int): return 2 * x')), - dict(id='2', kind='Code', name='c3', args=dict(code='def test(x: int): return 3 * x')), - dict(id='3', kind='JoinFormatter', name='join', args=dict(type='sum')) - ] - edges = [ - dict(iid='__start__', oid='0'), - dict(iid='__start__', oid='1'), - dict(iid='__start__', oid='2'), - dict(iid='0', oid='3'), - dict(iid='1', oid='3'), - dict(iid='2', oid='3'), - dict(iid='3', oid='__end__') - ] + nodes = [dict(id='0', kind='Code', name='c1', args=dict(code='def test(x: int): return x')), + dict(id='1', kind='Code', name='c2', args=dict(code='def test(x: int): return 2 * x')), + dict(id='2', kind='Code', name='c3', args=dict(code='def test(x: int): return 3 * x')), + dict(id='3', kind='JoinFormatter', name='join', args=dict(type='sum'))] + edges = [dict(iid='__start__', oid='0'), dict(iid='__start__', oid='1'), dict(iid='__start__', oid='2'), + dict(iid='0', oid='3'), dict(iid='1', oid='3'), dict(iid='2', oid='3'), dict(iid='3', oid='__end__')] gid = engine.start(nodes, edges) assert engine.run(gid, 1) == 6 assert engine.run(gid, '1') == '111111' assert engine.run(gid, [1]) == [1, 1, 1, 1, 1, 1] def test_engine_join_todict(self): - nodes = [ - dict(id='0', kind='Code', name='c1', args=dict(code='def test(x: int): return x')), - dict(id='1', kind='Code', name='c2', args=dict(code='def test(x: int): return 2 * x')), - dict(id='2', kind='Code', name='c3', args=dict(code='def test(x: int): return 3 * x')), - dict(id='3', kind='JoinFormatter', name='join', args=dict(type='to_dict', names=['a', 'b', 'c'])) - ] - edges = [ - dict(iid='__start__', oid='0'), - dict(iid='__start__', oid='1'), - dict(iid='__start__', oid='2'), - dict(iid='0', oid='3'), - dict(iid='1', oid='3'), - dict(iid='2', oid='3'), - dict(iid='3', oid='__end__') - ] + nodes = [dict(id='0', kind='Code', name='c1', args=dict(code='def test(x: int): return x')), + dict(id='1', kind='Code', name='c2', args=dict(code='def test(x: int): return 2 * x')), + dict(id='2', kind='Code', name='c3', args=dict(code='def test(x: int): return 3 * x')), + dict(id='3', kind='JoinFormatter', name='join', args=dict(type='to_dict', names=['a', 'b', 'c']))] + edges = [dict(iid='__start__', oid='0'), dict(iid='__start__', oid='1'), dict(iid='__start__', oid='2'), + dict(iid='0', oid='3'), dict(iid='1', oid='3'), dict(iid='2', oid='3'), dict(iid='3', oid='__end__')] engine = LightEngine() gid = engine.start(nodes, edges) assert engine.run(gid, 1) == dict(a=1, b=2, c=3) @@ -418,10 +342,9 @@ def test_engine_update(self): plus1 = dict(id='1', kind='Code', name='m1', args=dict(code='def test(x: int):\n return 1 + x\n')) double = dict(id='2', kind='Code', name='m2', args=dict(code='def test(x: int):\n return 2 * x\n')) square = dict(id='3', kind='Code', name='m3', args=dict(code='def test(x: int):\n return x * x\n')) - ifs = dict(id='4', - kind='Ifs', - name='i1', - args=dict(cond='def cond(x): return x < 10', true=[plus1, double], false=[square])) + ifs = dict(id='4', kind='Ifs', name='i1', args=dict( + cond='def cond(x): return x < 10', true=[plus1, double], false=[square] + )) nodes = [ifs] edges = [dict(iid='__start__', oid='4'), dict(iid='4', oid='__end__')] engine = LightEngine() @@ -431,10 +354,9 @@ def test_engine_update(self): assert engine.run(gid, 10) == 100 double = dict(id='2', kind='Code', name='m2', args=dict(code='def test(x: int):\n return 3 * x\n')) - ifs = dict(id='4', - kind='Ifs', - name='i1', - args=dict(cond='def cond(x): return x < 10', true=[plus1, double], false=[square])) + ifs = dict(id='4', kind='Ifs', name='i1', args=dict( + cond='def cond(x): return x < 10', true=[plus1, double], false=[square] + )) nodes = [ifs] engine.update(gid, nodes, edges) @@ -443,21 +365,12 @@ def test_engine_update(self): assert engine.run(gid, 10) == 100 def test_engine_join_join(self): - nodes = [ - dict(id='0', kind='Code', name='c1', args=dict(code='def test(x: int): return x')), - dict(id='1', kind='Code', name='c2', args=dict(code='def test(x: int): return 2 * x')), - dict(id='2', kind='Code', name='c3', args=dict(code='def test(x: int): return 3 * x')), - dict(id='3', kind='JoinFormatter', name='join', args=dict(type='join')) - ] - edges = [ - dict(iid='__start__', oid='0'), - dict(iid='__start__', oid='1'), - dict(iid='__start__', oid='2'), - dict(iid='0', oid='3'), - dict(iid='1', oid='3'), - dict(iid='2', oid='3'), - dict(iid='3', oid='__end__') - ] + nodes = [dict(id='0', kind='Code', name='c1', args=dict(code='def test(x: int): return x')), + dict(id='1', kind='Code', name='c2', args=dict(code='def test(x: int): return 2 * x')), + dict(id='2', kind='Code', name='c3', args=dict(code='def test(x: int): return 3 * x')), + dict(id='3', kind='JoinFormatter', name='join', args=dict(type='join'))] + edges = [dict(iid='__start__', oid='0'), dict(iid='__start__', oid='1'), dict(iid='__start__', oid='2'), + dict(iid='0', oid='3'), dict(iid='1', oid='3'), dict(iid='2', oid='3'), dict(iid='3', oid='__end__')] engine = LightEngine() gid = engine.start(nodes, edges) assert engine.run(gid, '1') == '111111' @@ -469,10 +382,9 @@ def test_engine_join_join(self): def test_engine_server(self): nodes = [dict(id='1', kind='Code', name='m1', args=dict(code='def test(x: int):\n return 2 * x\n'))] edges = [dict(iid='__start__', oid='1'), dict(iid='1', oid='__end__')] - resources = [ - dict(id='2', kind='server', name='s1', args=dict(port=None)), - dict(id='3', kind='web', name='w1', args=dict(port=None, title='网页', history=[], audio=False)) - ] + resources = [dict(id='2', kind='server', name='s1', args=dict(port=None)), + dict(id='3', kind='web', name='w1', args=dict(port=None, title='网页', history=[], audio=False)) + ] engine = LightEngine() gid = engine.start(nodes, edges, resources, gid='graph-1') assert engine.status(gid) == {'1': 'running', '2': lazyllm.launcher.Status.Running, '3': 'running'} @@ -528,24 +440,11 @@ def test_engine_httptool(self): dict(id='0', kind='Code', name='code1', args=dict(code='def p1(): return "foo"')), dict(id='1', kind='Code', name='code2', args=dict(code='def p2(): return "bar"')), dict(id='2', kind='Code', name='code3', args=dict(code='def h1(): return "baz"')), - dict(id='3', - kind='HttpTool', - name='http', - args=dict(method='GET', - url=url, - params=params, - headers=headers, - _lazyllm_arg_names=['p1', 'p2', 'h1'])) - ] - edges = [ - dict(iid='__start__', oid='0'), - dict(iid='__start__', oid='1'), - dict(iid='__start__', oid='2'), - dict(iid='0', oid='3'), - dict(iid='1', oid='3'), - dict(iid='2', oid='3'), - dict(iid='3', oid='__end__') + dict(id='3', kind='HttpTool', name='http', args=dict( + method='GET', url=url, params=params, headers=headers, _lazyllm_arg_names=['p1', 'p2', 'h1'])) ] + edges = [dict(iid='__start__', oid='0'), dict(iid='__start__', oid='1'), dict(iid='__start__', oid='2'), + dict(iid='0', oid='3'), dict(iid='1', oid='3'), dict(iid='2', oid='3'), dict(iid='3', oid='__end__')] engine = LightEngine() gid = engine.start(nodes, edges, gid='graph-1') @@ -564,24 +463,13 @@ def test_engine_status(self): square = dict(id='4', kind='Code', name='m3', args=dict(code='def test(x: int):\n return x * x\n')) subgraph = dict(id='5', kind='SubGraph', name='subgraph', args=dict(nodes=[double, plus1])) - ifs = dict(id='6', - kind='Ifs', - name='i1', - args=dict(cond='def cond(x): return x % 2 == 0', true=plus1, false=[square])) - loop = dict(id='7', - kind='Loop', - name='loop', - args=dict(stop_condition='def cond(x): return x > 8', nodes=[double])) - - switch = dict(id='8', - kind='Switch', - name='sw1', - args=dict(judge_on_full_input=True, nodes={ - 1: [plus1, subgraph], - 2: ifs, - 3: loop, - 5: [ifs] - })) + ifs = dict(id='6', kind='Ifs', name='i1', args=dict( + cond='def cond(x): return x % 2 == 0', true=plus1, false=[square])) + loop = dict(id='7', kind='Loop', name='loop', args=dict( + stop_condition='def cond(x): return x > 8', nodes=[double])) + + switch = dict(id='8', kind='Switch', name='sw1', args=dict(judge_on_full_input=True, nodes={ + 1: [plus1, subgraph], 2: ifs, 3: loop, 5: [ifs]})) warp = dict(id='9', kind='Warp', name='w1', args=dict(nodes=[switch, plus1])) join = dict(id='10', kind='JoinFormatter', name='join', args=dict(type='join', symbol=', ')) @@ -590,28 +478,14 @@ def test_engine_status(self): gid = engine.start(nodes, [], resources) assert '6, 4, 13, 26' in engine.run(gid, 1, 2, 3, 5) - assert engine.status(gid) == { - '9': { - '8': { - '2': 'running', - '5': { - '3': 'running', - '2': 'running' - }, - '6': { - '2': 'running', - '4': 'running' - }, - '7': { - '3': 'running' - } - }, - '2': 'running' - }, - '10': 'running', - '1': 'running', - '0': lazyllm.launcher.Status.Running - } + assert engine.status(gid) == {'9': {'8': {'2': 'running', + '5': {'3': 'running', '2': 'running'}, + '6': {'2': 'running', '4': 'running'}, + '7': {'3': 'running'}}, + '2': 'running'}, + '10': 'running', + '1': 'running', + '0': lazyllm.launcher.Status.Running} class TestEngineRAG(object): @@ -619,26 +493,15 @@ class TestEngineRAG(object): def test_rag(self): resources = [ dict(id='00', kind='LocalEmbedding', name='e1', args=dict(base_model='bge-large-zh-v1.5')), - dict(id='0', kind='Document', name='d1', args=dict(dataset_path='rag_master', embed='00')) - ] - nodes = [ - dict(id='1', - kind='Retriever', - name='ret1', - args=dict(doc='0', group_name='CoarseChunk', similarity='bm25_chinese', topk=3)), - dict(id='4', - kind='Reranker', - name='rek1', - args=dict(type='ModuleReranker', - output_format='content', - join=True, - arguments=dict(model="bge-reranker-large", topk=1))), - dict(id='5', - kind='Code', - name='c1', - args=dict(code='def test(nodes, query): return f\'context_str={nodes}, query={query}\'')), - dict(id='6', kind='LocalLLM', name='m1', args=dict(base_model='', deploy_method='dummy')) - ] + dict(id='0', kind='Document', name='d1', args=dict(dataset_path='rag_master', embed='00'))] + nodes = [dict(id='1', kind='Retriever', name='ret1', + args=dict(doc='0', group_name='CoarseChunk', similarity='bm25_chinese', topk=3)), + dict(id='4', kind='Reranker', name='rek1', + args=dict(type='ModuleReranker', output_format='content', join=True, + arguments=dict(model="bge-reranker-large", topk=1))), + dict(id='5', kind='Code', name='c1', + args=dict(code='def test(nodes, query): return f\'context_str={nodes}, query={query}\'')), + dict(id='6', kind='LocalLLM', name='m1', args=dict(base_model='', deploy_method='dummy'))] edges = [dict(iid='__start__', oid='1'), dict(iid='1', oid='4'), dict(iid='__start__', oid='4'), dict(iid='4', oid='5'), dict(iid='__start__', oid='5'), dict(iid='5', oid='6'), dict(iid='6', oid='__end__')] @@ -661,3 +524,4 @@ def test_rag(self): engine = LightEngine() engine.update(gid, nodes, edges, resources) assert '观天之道，执天之行' in engine.run(gid, '何为天道?') + From 98c28e311595d94da4ffffc831fbe302e7c6afe9 Mon Sep 17 00:00:00 2001 From: zhangyongchao Date: Tue, 12 Nov 2024 17:23:06 +0800 Subject: [PATCH 04/12] add testcase for code_block data_reflow --- tests/basic_tests/test_engine.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/basic_tests/test_engine.py b/tests/basic_tests/test_engine.py index 167e9ef2..8f390750 100644 --- a/tests/basic_tests/test_engine.py +++ b/tests/basic_tests/test_engine.py @@ -71,7 +71,8 @@ def test_engine_code(self): def test_engine_switch(self): plus1 = dict(id='1', kind='Code', name='m1', args=dict(code='def test(x: int):\n return 1 + x\n')) double = dict(id='2', kind='Code', name='m2', args=dict(code='def test(x: int):\n return 2 * x\n')) - square = dict(id='3', kind='Code', name='m3', args=dict(code='def test(x: int):\n return x * x\n')) + square = dict(id='3', kind='Code', name='m3', + args=dict(code='def test(x: int):\n return x * x\n', _lazyllm_enable_report=True)) switch = dict( id="4", kind="Switch", From a69ade50cf28d6aff21e1e8a08b3ddeac06ac541 Mon Sep 17 00:00:00 2001 From: zhangyongchao Date: Tue, 12 Nov 2024 17:26:37 +0800 Subject: [PATCH 05/12] fix lint problem --- tests/basic_tests/test_engine.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/basic_tests/test_engine.py b/tests/basic_tests/test_engine.py index 8f390750..6f7f2d0b 100644 --- a/tests/basic_tests/test_engine.py +++ b/tests/basic_tests/test_engine.py @@ -255,7 +255,8 @@ def test_engine_formatter_end(self): dict(id='2', kind='Code', name='m2', args=dict(code='def test1(x: int):\n return [[x, 2*x], [3*x, 4*x]]\n')), # two unused node - dict(id='3', kind='Code', name='m3', args=dict(code='def test2(x: int):\n return dict(a=1, b=x * x)\n')), + dict(id='3', kind='Code', name='m3', + args=dict(code='def test2(x: int):\n return dict(a=1, b=x * x)\n')), dict(id='4', kind='Code', name='m4', args=dict(code='def test3(x, y, z):\n return f"{x}{y}{z}"\n'))] edges = [dict(iid='__start__', oid='1'), dict(iid='__start__', oid='2'), dict(iid='2', oid='__end__'), dict(iid='1', oid='__end__')] @@ -525,4 +526,3 @@ def test_rag(self): engine = LightEngine() engine.update(gid, nodes, edges, resources) assert '观天之道，执天之行' in engine.run(gid, '何为天道?') - From be1f4df7a7f2c192a13aa3eaacbb0a8eedd5bd2e Mon Sep 17 00:00:00 2001 From: zhangyongchao Date: Tue, 12 Nov 2024 17:29:08 +0800 Subject: [PATCH 06/12] fix lint problem1 --- tests/basic_tests/test_engine.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/basic_tests/test_engine.py b/tests/basic_tests/test_engine.py index 6f7f2d0b..ac8d1b02 100644 --- a/tests/basic_tests/test_engine.py +++ b/tests/basic_tests/test_engine.py @@ -255,7 +255,7 @@ def test_engine_formatter_end(self): dict(id='2', kind='Code', name='m2', args=dict(code='def test1(x: int):\n return [[x, 2*x], [3*x, 4*x]]\n')), # two unused node - dict(id='3', kind='Code', name='m3', + dict(id='3', kind='Code', name='m3', args=dict(code='def test2(x: int):\n return dict(a=1, b=x * x)\n')), dict(id='4', kind='Code', name='m4', args=dict(code='def test3(x, y, z):\n return f"{x}{y}{z}"\n'))] edges = [dict(iid='__start__', oid='1'), dict(iid='__start__', oid='2'), dict(iid='2', oid='__end__'), From 2691ca565986a689e8c183511339511ff5aeaa5d Mon Sep 17 00:00:00 2001 From: zhangyongchao Date: Wed, 13 Nov 2024 17:15:39 +0800 Subject: [PATCH 07/12] 1. Fix bug: When hook is in server node, the report_url becomes empty 2. Add fastapi server for repot_url --- lazyllm/engine/engine.py | 5 +- lazyllm/engine/node_meta_hook.py | 13 ++-- tests/basic_tests/test_engine.py | 128 ++++++++++++++++++++++++++----- 3 files changed, 118 insertions(+), 28 deletions(-) diff --git a/lazyllm/engine/engine.py b/lazyllm/engine/engine.py index 1a3e54fe..bfd3dc12 100644 --- a/lazyllm/engine/engine.py +++ b/lazyllm/engine/engine.py @@ -12,6 +12,7 @@ # Each session will have a separate engine class Engine(object): __default_engine__ = None + REPORT_URL = "" def __init__(self): self._nodes = {'__start__': Node(id='__start__', kind='__start__', name='__start__'), @@ -55,7 +56,7 @@ def build_node(self, node) -> Callable: return _constructor.build(node) def set_report_url(self, url) -> None: - NodeMetaHook.URL = url + Engine.REPORT_URL = url def reset(self): for node in self._nodes: @@ -139,7 +140,7 @@ def _process_hook(self, node, module): NodeMetaHook.MODULEID_TO_WIDGETID[module._flow_id] = node.id else: return - node.func.register_hook(NodeMetaHook) + node.func.register_hook(NodeMetaHook(node.func, Engine.REPORT_URL, node.id)) _constructor = NodeConstructor() diff --git a/lazyllm/engine/node_meta_hook.py b/lazyllm/engine/node_meta_hook.py index ae1165c2..b61a10d8 100644 --- a/lazyllm/engine/node_meta_hook.py +++ b/lazyllm/engine/node_meta_hook.py @@ -19,7 +19,7 @@ class NodeMetaHook(LazyLLMHook): URL = "" MODULEID_TO_WIDGETID = {} - def __init__(self, obj): + def __init__(self, obj, url, front_id): if isinstance(obj, lazyllm.ModuleBase): self._uniqueid = obj._module_id elif isinstance(obj, lazyllm.FlowBase): @@ -28,16 +28,18 @@ def __init__(self, obj): raise TypeError(f"Expected 'obj' to be type of ModuleBase or FlowBase, but got {type(obj)}") self._meta_info = { MetaKeys.ID: str(self._uniqueid), - MetaKeys.SESSIONID: lazyllm.globals._sid, MetaKeys.TIMECOST: 0.0, MetaKeys.PROMPT_TOKENS: 0, MetaKeys.COMPLETION_TOKENS: 0, MetaKeys.INPUT: "", MetaKeys.OUTPUT: "", } + self._front_id = front_id + self._url = url def pre_hook(self, *args, **kwargs): arguments = {} + self._meta_info[MetaKeys.SESSIONID] = lazyllm.globals._sid if len(args) == 1: if isinstance(args[0], lazyllm.package) and len(args[0]) == 1: self._meta_info[MetaKeys.INPUT] = str(args[0][0]) @@ -55,8 +57,7 @@ def post_hook(self, output): if self._uniqueid in globals["usage"]: self._meta_info.update(globals["usage"]) - if self._meta_info[MetaKeys.ID] in self.MODULEID_TO_WIDGETID: - self._meta_info[MetaKeys.ID] = self.MODULEID_TO_WIDGETID[self._meta_info[MetaKeys.ID]] + self._meta_info[MetaKeys.ID] = self._front_id self._meta_info[MetaKeys.TIMECOST] = time.time() - self._meta_info[MetaKeys.TIMECOST] def report(self): @@ -64,6 +65,6 @@ def report(self): json_data = json.dumps(self._meta_info, ensure_ascii=False) try: lazyllm.LOG.info(f"meta_info: {self._meta_info}") - requests.post(self.URL, data=json_data, headers=headers) + requests.post(self._url, data=json_data, headers=headers) except Exception as e: - lazyllm.LOG.warning(f"Error sending collected data: {e}") + lazyllm.LOG.warning(f"Error sending collected data: {e}. URL: {self._url}") diff --git a/tests/basic_tests/test_engine.py b/tests/basic_tests/test_engine.py index ac8d1b02..9e9865fb 100644 --- a/tests/basic_tests/test_engine.py +++ b/tests/basic_tests/test_engine.py @@ -10,33 +10,60 @@ from fastapi import FastAPI from fastapi.responses import JSONResponse from fastapi.testclient import TestClient +import subprocess +import socket +import threading + +HOOK_PORT = 33733 +HOOK_ROUTE = "mock_post" +fastapi_code = """ +from fastapi import FastAPI +from fastapi.responses import JSONResponse +from fastapi.testclient import TestClient app = FastAPI() -@app.post("/mock_post") +@app.post("/{route}") async def receive_json(data: dict): + print("Received json data:", data) return JSONResponse(content=data) +if __name__ == "__main__": + import uvicorn + uvicorn.run(app, host="0.0.0.0", port={port}) +""".format( + port=HOOK_PORT, route=HOOK_ROUTE +) + + class TestEngine(unittest.TestCase): @classmethod def setUpClass(cls): - client = TestClient(app) - - def mock_report(self): - headers = {"Content-Type": "application/json; charset=utf-8"} - json_data = json.dumps(self._meta_info, ensure_ascii=False) - try: - lazyllm.LOG.info(f"meta_info: {self._meta_info}") - response = client.post(self.URL, data=json_data, headers=headers) - assert ( - response.json() == self._meta_info - ), "mock response should be same as input" - except Exception as e: - lazyllm.LOG.warning(f"Error sending collected data: {e}") - - NodeMetaHook.report = mock_report + cls.fastapi_process = subprocess.Popen( + ["python", "-c", fastapi_code], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + hostname = socket.gethostname() + ip_address = socket.gethostbyname(hostname) + cls.report_url = f"http://{ip_address}:{HOOK_PORT}/{HOOK_ROUTE}" + + def read_stdout(process): + for line in iter(process.stdout.readline, b''): + print("FastAPI Server Output: ", line.decode(), end='') + + cls.report_print_thread = threading.Thread( + target=read_stdout, args=(cls.fastapi_process,) + ) + cls.report_print_thread.daemon = True + cls.report_print_thread.start() + + @classmethod + def tearDownClass(cls): + cls.fastapi_process.terminate() + cls.fastapi_process.wait() @pytest.fixture(autouse=True) def run_around_tests(self): @@ -86,7 +113,7 @@ def test_engine_switch(self): nodes = [switch] edges = [dict(iid='__start__', oid='4'), dict(iid='4', oid='__end__')] engine = LightEngine() - engine.set_report_url("mock_post") + engine.set_report_url(self.report_url) gid = engine.start(nodes, edges) assert engine.run(gid, 1) == 2 assert engine.run(gid, 2) == 6 @@ -130,12 +157,73 @@ def test_engine_ifs(self): nodes = [ifs] edges = [dict(iid='__start__', oid='4'), dict(iid='4', oid='__end__')] engine = LightEngine() - engine.set_report_url("mock_post") + engine.set_report_url(self.report_url) gid = engine.start(nodes, edges) assert engine.run(gid, 1) == 4 assert engine.run(gid, 5) == 12 assert engine.run(gid, 10) == 100 + def test_data_reflow_in_server(self): + nodes = [ + { + "id": "debug-48f9e95a-e54c-4812-a41b-f61f635816b8-1731324014737", + "kind": "Code", + "name": "1731324014737", + "args": { + "code": "def main(x): return int(x) + 1", + "_lazyllm_enable_report": True, + }, + }, + { + "id": "debug-48f9e95a-e54c-4812-a41b-f61f635816b8-1731324018968", + "kind": "Code", + "name": "1731324018968", + "args": { + "code": "def main(x): return int(x) + 2", + "_lazyllm_enable_report": True, + }, + }, + { + "id": "debug-48f9e95a-e54c-4812-a41b-f61f635816b8-1731324021877", + "kind": "Code", + "name": "1731324021877", + "args": { + "code": "def main(x): return int(x) + 3", + "_lazyllm_enable_report": True, + }, + }, + ] + edges = [ + { + "iid": "__start__", + "oid": "debug-48f9e95a-e54c-4812-a41b-f61f635816b8-1731324014737", + }, + { + "iid": "debug-48f9e95a-e54c-4812-a41b-f61f635816b8-1731324014737", + "oid": "debug-48f9e95a-e54c-4812-a41b-f61f635816b8-1731324018968", + }, + { + "iid": "debug-48f9e95a-e54c-4812-a41b-f61f635816b8-1731324018968", + "oid": "debug-48f9e95a-e54c-4812-a41b-f61f635816b8-1731324021877", + }, + { + "iid": "debug-48f9e95a-e54c-4812-a41b-f61f635816b8-1731324021877", + "oid": "__end__", + }, + ] + resources = [ + { + "id": "adc5174c-3220-4b98-92a6-3897e9af8898", + "kind": "server", + "name": "adc5174c-3220-4b98-92a6-3897e9af8898", + "args": {}, + } + ] + engine = LightEngine() + engine.set_report_url(self.report_url) + gid = engine.start(nodes, edges, resources) + assert engine.run(gid, 1) == 7 + def test_engine_loop(self): nodes = [dict(id='1', kind='Code', name='code', args=dict(code='def square(x: int): return x * x'))] edges = [dict(iid='__start__', oid='1'), dict(iid='1', oid='__end__')] @@ -156,7 +244,7 @@ def test_engine_loop(self): edges = [dict(iid='__start__', oid='2'), dict(iid='2', oid='__end__')] engine = LightEngine() - engine.set_report_url("mock_post") + engine.set_report_url(self.report_url) gid = engine.start(nodes, edges) assert engine.run(gid, 2) == 16 @@ -179,7 +267,7 @@ def test_engine_warp(self): edges = [dict(iid='__start__', oid='2'), dict(iid='2', oid='__end__')] engine = LightEngine() - engine.set_report_url("mock_post") + engine.set_report_url(self.report_url) gid = engine.start(nodes, edges) assert engine.run(gid, 2, 3, 4, 5) == (4, 9, 16, 25) From d5c3dc70052d513a1f18718266cf53d6fd9e3b63 Mon Sep 17 00:00:00 2001 From: zhangyongchao Date: Wed, 13 Nov 2024 18:47:01 +0800 Subject: [PATCH 08/12] refactor and fix token usage key missing --- lazyllm/engine/engine.py | 16 ++-------------- lazyllm/engine/node_meta_hook.py | 5 +---- tests/basic_tests/test_engine.py | 2 ++ 3 files changed, 5 insertions(+), 18 deletions(-) diff --git a/lazyllm/engine/engine.py b/lazyllm/engine/engine.py index 85c86db5..831f7da4 100644 --- a/lazyllm/engine/engine.py +++ b/lazyllm/engine/engine.py @@ -85,7 +85,6 @@ def _impl(nid, recursive): yield id if recursive: yield from self.subnodes(id, True) - return list(_impl(nodeid, recursive)) @@ -105,7 +104,6 @@ def impl(f): for name in names: cls.builder_methods[name] = (f, subitems) return f - return impl # build node recursively @@ -171,11 +169,7 @@ def get_args(cls, key, value, builder_key=None): def _process_hook(self, node, module): if not node.enable_data_reflow: return - if isinstance(module, lazyllm.ModuleBase): - NodeMetaHook.MODULEID_TO_WIDGETID[module._module_id] = node.id - elif isinstance(module, lazyllm.LazyLLMFlowsBase): - NodeMetaHook.MODULEID_TO_WIDGETID[module._flow_id] = node.id - else: + if not isinstance(module, (lazyllm.ModuleBase, lazyllm.LazyLLMFlowsBase)): return node.func.register_hook(NodeMetaHook(node.func, Engine.REPORT_URL, node.id)) @@ -487,27 +481,22 @@ def _parse_py_data_by_formatter(self, data): else: raise TypeError('type should be one of sum/stack/to_dict/join') - @NodeConstructor.register('JoinFormatter') def make_join_formatter(type='sum', names=None, symbol=None): if type == 'file': return make_formatter('file', rule='merge') return JoinFormatter(type, names=names, symbol=symbol) - @NodeConstructor.register('Formatter') def make_formatter(ftype, rule): return getattr(lazyllm.formatter, ftype)(formatter=rule) - def return_a_wrapper_func(func): @functools.wraps(func) def wrapper_func(*args, **kwargs): return func(*args, **kwargs) - return wrapper_func - def _get_tools(tools): callable_list = [] for rid in tools: # `tools` is a list of ids in engine's resources @@ -520,12 +509,10 @@ def _get_tools(tools): callable_list.append(wrapper_func) return callable_list - @NodeConstructor.register('ToolsForLLM') def make_tools_for_llm(tools: List[str]): return lazyllm.tools.ToolManager(_get_tools(tools)) - @NodeConstructor.register('FunctionCall') def make_fc(llm: str, tools: List[str], algorithm: Optional[str] = None): f = ( @@ -566,6 +553,7 @@ def make_http_tool( class VQA(lazyllm.Module): + def __init__( self, base_model: Union[str, lazyllm.TrainableModule], diff --git a/lazyllm/engine/node_meta_hook.py b/lazyllm/engine/node_meta_hook.py index b61a10d8..dcab0632 100644 --- a/lazyllm/engine/node_meta_hook.py +++ b/lazyllm/engine/node_meta_hook.py @@ -16,8 +16,6 @@ class MetaKeys: class NodeMetaHook(LazyLLMHook): - URL = "" - MODULEID_TO_WIDGETID = {} def __init__(self, obj, url, front_id): if isinstance(obj, lazyllm.ModuleBase): @@ -56,7 +54,7 @@ def post_hook(self, output): self._meta_info[MetaKeys.OUTPUT] = str(output) if self._uniqueid in globals["usage"]: - self._meta_info.update(globals["usage"]) + self._meta_info.update(globals["usage"][self._uniqueid]) self._meta_info[MetaKeys.ID] = self._front_id self._meta_info[MetaKeys.TIMECOST] = time.time() - self._meta_info[MetaKeys.TIMECOST] @@ -64,7 +62,6 @@ def report(self): headers = {"Content-Type": "application/json; charset=utf-8"} json_data = json.dumps(self._meta_info, ensure_ascii=False) try: - lazyllm.LOG.info(f"meta_info: {self._meta_info}") requests.post(self._url, data=json_data, headers=headers) except Exception as e: lazyllm.LOG.warning(f"Error sending collected data: {e}. URL: {self._url}") diff --git a/tests/basic_tests/test_engine.py b/tests/basic_tests/test_engine.py index f2fe10ae..70c680eb 100644 --- a/tests/basic_tests/test_engine.py +++ b/tests/basic_tests/test_engine.py @@ -27,6 +27,7 @@ @app.post("/{route}") async def receive_json(data: dict): print("Received json data:", data) + assert "prompt_tokens" in data return JSONResponse(content=data) @@ -62,6 +63,7 @@ def read_stdout(process): @classmethod def tearDownClass(cls): + time.sleep(3) cls.fastapi_process.terminate() cls.fastapi_process.wait() From 64eaad8a2e53b7fdcf805c767c36d53efb3a621d Mon Sep 17 00:00:00 2001 From: zhangyongchao Date: Wed, 13 Nov 2024 18:54:24 +0800 Subject: [PATCH 09/12] restore unchanged part to original format --- lazyllm/engine/engine.py | 347 +++++++++++---------------------------- 1 file changed, 94 insertions(+), 253 deletions(-) diff --git a/lazyllm/engine/engine.py b/lazyllm/engine/engine.py index 831f7da4..5a0b4553 100644 --- a/lazyllm/engine/engine.py +++ b/lazyllm/engine/engine.py @@ -15,10 +15,8 @@ class Engine(object): REPORT_URL = "" def __init__(self): - self._nodes = { - '__start__': Node(id='__start__', kind='__start__', name='__start__'), - '__end__': Node(id='__end__', kind='__end__', name='__end__'), - } + self._nodes = {'__start__': Node(id='__start__', kind='__start__', name='__start__'), + '__end__': Node(id='__end__', kind='__end__', name='__end__')} def __new__(cls): if cls is not Engine: @@ -30,38 +28,29 @@ def set_default(cls, engine: Type): cls.__default_engine__ = engine @overload - def start(self, nodes: str) -> None: ... + def start(self, nodes: str) -> None: + ... @overload - def start(self, nodes: Dict[str, Any]) -> None: ... + def start(self, nodes: Dict[str, Any]) -> None: + ... @overload - def start( - self, - nodes: List[Dict] = [], - edges: List[Dict] = [], - resources: List[Dict] = [], - gid: Optional[str] = None, - name: Optional[str] = None, - ) -> str: ... + def start(self, nodes: List[Dict] = [], edges: List[Dict] = [], resources: List[Dict] = [], + gid: Optional[str] = None, name: Optional[str] = None) -> str: + ... @overload - def update(self, nodes: List[Dict]) -> None: ... + def update(self, nodes: List[Dict]) -> None: + ... @overload - def update( - self, - gid: str, - nodes: List[Dict], - edges: List[Dict] = [], - resources: List[Dict] = [], - ) -> str: ... + def update(self, gid: str, nodes: List[Dict], edges: List[Dict] = [], + resources: List[Dict] = []) -> str: + ... - def release_node(self, nodeid: str): - pass - - def stop(self, node_id: Optional[str] = None, task_name: Optional[str] = None): - pass + def release_node(self, nodeid: str): pass + def stop(self, node_id: Optional[str] = None, task_name: Optional[str] = None): pass def build_node(self, node) -> Node: return _constructor.build(node) @@ -83,8 +72,7 @@ def subnodes(self, nodeid: str, recursive: bool = False): def _impl(nid, recursive): for id in self._nodes[nid].subitems: yield id - if recursive: - yield from self.subnodes(id, True) + if recursive: yield from self.subnodes(id, True) return list(_impl(nodeid, recursive)) @@ -92,13 +80,8 @@ class NodeConstructor(object): builder_methods = dict() @classmethod - def register( - cls, - *names: Union[List[str], str], - subitems: Optional[Union[str, List[str]]] = None, - ): - if len(names) == 1 and isinstance(names[0], (tuple, list)): - names = names[0] + def register(cls, *names: Union[List[str], str], subitems: Optional[Union[str, List[str]]] = None): + if len(names) == 1 and isinstance(names[0], (tuple, list)): names = names[0] def impl(f): for name in names: @@ -110,26 +93,13 @@ def impl(f): def build(self, node: Node): if node.kind.startswith('__') and node.kind.endswith('__'): return None - node.arg_names = ( - node.args.pop('_lazyllm_arg_names', None) - if isinstance(node.args, dict) - else None - ) - node.enable_data_reflow = ( - node.args.pop('_lazyllm_enable_report', False) - if isinstance(node.args, dict) - else False - ) + node.arg_names = node.args.pop('_lazyllm_arg_names', None) if isinstance(node.args, dict) else None + node.enable_data_reflow = (node.args.pop('_lazyllm_enable_report', False) + if isinstance(node.args, dict) else False) if node.kind in NodeConstructor.builder_methods: createf, node.subitem_name = NodeConstructor.builder_methods[node.kind] - node.func = ( - createf(**node.args) - if isinstance(node.args, dict) - and set(node.args.keys()).issubset( - set(inspect.getfullargspec(createf).args) - ) - else createf(node.args) - ) + node.func = createf(**node.args) if isinstance(node.args, dict) and set(node.args.keys()).issubset( + set(inspect.getfullargspec(createf).args)) else createf(node.args) self._process_hook(node, node.func) return node @@ -137,9 +107,7 @@ def build(self, node: Node): init_args, build_args, other_args = dict(), dict(), dict() def get_args(cls, key, value, builder_key=None): - node_args = ( - node_msgs[cls][builder_key][key] if builder_key else node_msgs[cls][key] - ) + node_args = node_msgs[cls][builder_key][key] if builder_key else node_msgs[cls][key] if node_args.type == Node: return Engine().build_node(value).func return node_args.getattr_f(value) if node_args.getattr_f else value @@ -151,11 +119,8 @@ def get_args(cls, key, value, builder_key=None): build_args[key] = get_args('builder_argument', key, value) elif '.' in key: builder_key, key = key.split('.') - if builder_key not in other_args: - other_args[builder_key] = dict() - other_args[builder_key][key] = get_args( - 'other_arguments', key, value, builder_key=builder_key - ) + if builder_key not in other_args: other_args[builder_key] = dict() + other_args[builder_key][key] = get_args('other_arguments', key, value, builder_key=builder_key) else: raise KeyError(f'Invalid key `{key}` found') @@ -182,20 +147,12 @@ def __init__(self, g: lazyllm.graph, server: Node, web: Node): super().__init__() self._g = lazyllm.ActionModule(g) if server: - if server.args.get('port'): - raise NotImplementedError('Port is not supported now') + if server.args.get('port'): raise NotImplementedError('Port is not supported now') self._g = lazyllm.ServerModule(g) if web: port = self._get_port(web.args['port']) - self._web = lazyllm.WebModule( - g, - port=port, - title=web.args['title'], - audio=web.args['audio'], - history=[ - Engine().build_node(h).func for h in web.args.get('history', []) - ], - ) + self._web = lazyllm.WebModule(g, port=port, title=web.args['title'], audio=web.args['audio'], + history=[Engine().build_node(h).func for h in web.args.get('history', [])]) def forward(self, *args, **kw): return self._g(*args, **kw) @@ -203,13 +160,11 @@ def forward(self, *args, **kw): # TODO(wangzhihong) def _update(self, *, mode=None, recursive=True): super(__class__, self)._update(mode=mode, recursive=recursive) - if hasattr(self, '_web'): - self._web.start() + if hasattr(self, '_web'): self._web.start() return self def _get_port(self, port): - if not port: - return None + if not port: return None elif ',' in port: return list(int(p.strip()) for p in port.split(',')) elif '-' in port: @@ -241,9 +196,7 @@ def __init__(self, graph: ServerGraph, kind: str, args: Dict): self._args = args def status(self): - return ( - self._graph._g.status if self._kind == 'server' else self._graph._web.status - ) + return self._graph._g.status if self._kind == 'server' else self._graph._web.status @NodeConstructor.register('web', 'server') @@ -252,34 +205,18 @@ def make_server_resource(kind: str, graph: ServerGraph, args: Dict[str, Any]): @NodeConstructor.register('Graph', 'SubGraph', subitems=['nodes', 'resources']) -def make_graph( - nodes: List[dict], - edges: List[Union[List[str], dict]] = [], - resources: List[dict] = [], - enable_server=True, -): +def make_graph(nodes: List[dict], edges: List[Union[List[str], dict]] = [], + resources: List[dict] = [], enable_server=True): engine = Engine() server_resources = dict(server=None, web=None) for resource in resources: if resource['kind'] in server_resources: - assert ( - enable_server - ), 'Web and Api server are not allowed outside graph and subgraph' - assert ( - server_resources[resource['kind']] is None - ), f'Duplicated {resource["kind"]} resource' - server_resources[resource['kind']] = Node( - id=resource['id'], - kind=resource['kind'], - name=resource['name'], - args=resource['args'], - ) - - resources = [ - engine.build_node(resource) - for resource in resources - if resource['kind'] not in server_resources - ] + assert enable_server, 'Web and Api server are not allowed outside graph and subgraph' + assert server_resources[resource['kind']] is None, f'Duplicated {resource["kind"]} resource' + server_resources[resource['kind']] = Node(id=resource['id'], kind=resource['kind'], + name=resource['name'], args=resource['args']) + + resources = [engine.build_node(resource) for resource in resources if resource['kind'] not in server_resources] nodes: List[Node] = [engine.build_node(node) for node in nodes] with graph() as g: @@ -288,31 +225,19 @@ def make_graph( g.set_node_arg_name([node.arg_names for node in nodes]) if not edges: - edges = ( - [dict(iid='__start__', oid=nodes[0].id)] - + [ - dict(iid=nodes[i].id, oid=nodes[i + 1].id) - for i in range(len(nodes) - 1) - ] - + [dict(iid=nodes[-1].id, oid='__end__')] - ) + edges = ([dict(iid='__start__', oid=nodes[0].id)] + [ + dict(iid=nodes[i].id, oid=nodes[i + 1].id) for i in range(len(nodes) - 1)] + [ + dict(iid=nodes[-1].id, oid='__end__')]) for edge in edges: - if isinstance(edge, (tuple, list)): - edge = dict(iid=edge[0], oid=edge[1]) + if isinstance(edge, (tuple, list)): edge = dict(iid=edge[0], oid=edge[1]) if formatter := edge.get('formatter'): - assert formatter.startswith(('*[', '[', '}')) and formatter.endswith( - (']', '}') - ) + assert formatter.startswith(('*[', '[', '}')) and formatter.endswith((']', '}')) formatter = lazyllm.formatter.JsonLike(formatter) if 'constant' in edge: g.add_const_edge(edge['constant'], engine._nodes[edge['oid']].name) else: - g.add_edge( - engine._nodes[edge['iid']].name, - engine._nodes[edge['oid']].name, - formatter, - ) + g.add_edge(engine._nodes[edge['iid']].name, engine._nodes[edge['oid']].name, formatter) sg = ServerGraph(g, server_resources['server'], server_resources['web']) for kind, node in server_resources.items(): @@ -354,7 +279,7 @@ def _build_pipeline(nodes): def make_switch(judge_on_full_input: bool, nodes: Dict[str, List[dict]]): with switch(judge_on_full_input=judge_on_full_input) as sw: for cond, nodes in nodes.items(): - sw.case[cond :: _build_pipeline(nodes)] + sw.case[cond::_build_pipeline(nodes)] return sw @@ -364,93 +289,48 @@ def make_warp(nodes: List[dict], edges: List[dict] = [], resources: List[dict] = @NodeConstructor.register('Loop', subitems=['nodes', 'resources']) -def make_loop( - stop_condition: str, - nodes: List[dict], - edges: List[dict] = [], - resources: List[dict] = [], - judge_on_full_input: bool = True, -): +def make_loop(stop_condition: str, nodes: List[dict], edges: List[dict] = [], + resources: List[dict] = [], judge_on_full_input: bool = True): stop_condition = make_code(stop_condition) - return lazyllm.loop( - make_graph(nodes, edges, resources, enable_server=False), - stop_condition=stop_condition, - judge_on_full_input=judge_on_full_input, - ) + return lazyllm.loop(make_graph(nodes, edges, resources, enable_server=False), + stop_condition=stop_condition, judge_on_full_input=judge_on_full_input) @NodeConstructor.register('Ifs', subitems=['true', 'false']) -def make_ifs( - cond: str, true: List[dict], false: List[dict], judge_on_full_input: bool = True -): +def make_ifs(cond: str, true: List[dict], false: List[dict], judge_on_full_input: bool = True): assert judge_on_full_input, 'judge_on_full_input only support True now' - return lazyllm.ifs( - make_code(cond), tpath=_build_pipeline(true), fpath=_build_pipeline(false) - ) + return lazyllm.ifs(make_code(cond), tpath=_build_pipeline(true), fpath=_build_pipeline(false)) @NodeConstructor.register('Intention', subitems=['nodes:dict']) -def make_intention( - base_model: str, - nodes: Dict[str, List[dict]], - prompt: str = '', - constrain: str = '', - attention: str = '', -): - with IntentClassifier( - Engine().build_node(base_model).func, - prompt=prompt, - constrain=constrain, - attention=attention, - ) as ic: +def make_intention(base_model: str, nodes: Dict[str, List[dict]], + prompt: str = '', constrain: str = '', attention: str = ''): + with IntentClassifier(Engine().build_node(base_model).func, + prompt=prompt, constrain=constrain, attention=attention) as ic: for cond, nodes in nodes.items(): if isinstance(nodes, list) and len(nodes) > 1: f = pipeline([Engine().build_node(node).func for node in nodes]) else: - f = ( - Engine() - .build_node(nodes[0] if isinstance(nodes, list) else nodes) - .func - ) + f = Engine().build_node(nodes[0] if isinstance(nodes, list) else nodes).func ic.case[cond::f] return ic @NodeConstructor.register('Document') -def make_document( - dataset_path: str, - embed: Node = None, - create_ui: bool = False, - server: bool = False, - node_group: List = [], -): +def make_document(dataset_path: str, embed: Node = None, create_ui: bool = False, + server: bool = False, node_group: List = []): document = lazyllm.tools.rag.Document( - dataset_path, - Engine().build_node(embed).func if embed else None, - server=server, - manager=create_ui, - ) + dataset_path, Engine().build_node(embed).func if embed else None, server=server, manager=create_ui) for group in node_group: - if group['transform'] == 'LLMParser': - group['llm'] = Engine().build_node(group['llm']).func - elif group['transform'] == 'FuncNode': - group['function'] = make_code(group['function']) + if group['transform'] == 'LLMParser': group['llm'] = Engine().build_node(group['llm']).func + elif group['transform'] == 'FuncNode': group['function'] = make_code(group['function']) document.create_node_group(**group) return document - @NodeConstructor.register('Reranker') -def make_reranker( - type: str = 'ModuleReranker', - target: Optional[str] = None, - output_format: Optional[str] = None, - join: Union[bool, str] = False, - arguments: Dict = {}, -): - return lazyllm.tools.Reranker( - type, target=target, output_format=output_format, join=join, **arguments - ) - +def make_reranker(type: str = 'ModuleReranker', target: Optional[str] = None, + output_format: Optional[str] = None, join: Union[bool, str] = False, arguments: Dict = {}): + return lazyllm.tools.Reranker(type, target=target, output_format=output_format, join=join, **arguments) class JoinFormatter(lazyllm.components.FormatterBase): def __init__(self, type, *, names=None, symbol=None): @@ -461,17 +341,10 @@ def __init__(self, type, *, names=None, symbol=None): def _parse_py_data_by_formatter(self, data): if self.type == 'sum': assert len(data) > 0, 'Cannot sum empty inputs' - if isinstance(data[0], str): - return ''.join(data) + if isinstance(data[0], str): return ''.join(data) return sum(data, type(data[0])()) elif self.type == 'stack': - return ( - list(data) - if isinstance(data, package) - else [ - data, - ] - ) + return list(data) if isinstance(data, package) else [data,] elif self.type == 'to_dict': assert self.names and len(self.names) == len(data) return {k: v for k, v in zip(self.names, data)} @@ -483,8 +356,7 @@ def _parse_py_data_by_formatter(self, data): @NodeConstructor.register('JoinFormatter') def make_join_formatter(type='sum', names=None, symbol=None): - if type == 'file': - return make_formatter('file', rule='merge') + if type == 'file': return make_formatter('file', rule='merge') return JoinFormatter(type, names=names, symbol=symbol) @NodeConstructor.register('Formatter') @@ -515,56 +387,34 @@ def make_tools_for_llm(tools: List[str]): @NodeConstructor.register('FunctionCall') def make_fc(llm: str, tools: List[str], algorithm: Optional[str] = None): - f = ( - lazyllm.tools.PlanAndSolveAgent - if algorithm == 'PlanAndSolve' - else ( - lazyllm.tools.ReWOOAgent - if algorithm == 'ReWOO' - else ( - lazyllm.tools.ReactAgent - if algorithm == 'React' - else lazyllm.tools.FunctionCallAgent - ) - ) - ) + f = lazyllm.tools.PlanAndSolveAgent if algorithm == 'PlanAndSolve' else \ + lazyllm.tools.ReWOOAgent if algorithm == 'ReWOO' else \ + lazyllm.tools.ReactAgent if algorithm == 'React' else lazyllm.tools.FunctionCallAgent return f(Engine().build_node(llm).func, _get_tools(tools)) - @NodeConstructor.register('HttpTool') -def make_http_tool( - method: Optional[str] = None, - url: Optional[str] = None, - params: Optional[Dict[str, str]] = None, - headers: Optional[Dict[str, str]] = None, - body: Optional[str] = None, - timeout: int = 10, - proxies: Optional[Dict[str, str]] = None, - code_str: Optional[str] = None, - vars_for_code: Optional[Dict[str, Any]] = None, - doc: Optional[str] = None, -): - instance = lazyllm.tools.HttpTool( - method, url, params, headers, body, timeout, proxies, code_str, vars_for_code - ) +def make_http_tool(method: Optional[str] = None, + url: Optional[str] = None, + params: Optional[Dict[str, str]] = None, + headers: Optional[Dict[str, str]] = None, + body: Optional[str] = None, + timeout: int = 10, + proxies: Optional[Dict[str, str]] = None, + code_str: Optional[str] = None, + vars_for_code: Optional[Dict[str, Any]] = None, + doc: Optional[str] = None): + instance = lazyllm.tools.HttpTool(method, url, params, headers, body, timeout, proxies, + code_str, vars_for_code) if doc: instance.__doc__ = doc return instance class VQA(lazyllm.Module): - - def __init__( - self, - base_model: Union[str, lazyllm.TrainableModule], - file_resource_id: Optional[str], - ): + def __init__(self, base_model: Union[str, lazyllm.TrainableModule], file_resource_id: Optional[str]): super().__init__() - self.vqa = self._vqa = ( - lazyllm.TrainableModule(base_model).deploy_method(lazyllm.deploy.LMDeploy) - if not isinstance(base_model, lazyllm.TrainableModule) - else base_model - ) + self.vqa = self._vqa = (lazyllm.TrainableModule(base_model).deploy_method(lazyllm.deploy.LMDeploy) + if not isinstance(base_model, lazyllm.TrainableModule) else base_model) self._file_resource_id = file_resource_id if file_resource_id: with pipeline() as self.vqa: @@ -588,17 +438,10 @@ def make_vqa(base_model: str, file_resource_id: Optional[str] = None): @NodeConstructor.register('SharedLLM') -def make_shared_llm( - llm: str, prompt: Optional[str] = None, file_resource_id: Optional[str] = None -): +def make_shared_llm(llm: str, prompt: Optional[str] = None, file_resource_id: Optional[str] = None): llm = Engine().build_node(llm).func - if file_resource_id: - assert isinstance(llm, VQA), 'file_resource_id is only supported in VQA' - return ( - VQA(llm._vqa.share(prompt=prompt), file_resource_id) - if file_resource_id - else llm.share(prompt=prompt) - ) + if file_resource_id: assert isinstance(llm, VQA), 'file_resource_id is only supported in VQA' + return VQA(llm._vqa.share(prompt=prompt), file_resource_id) if file_resource_id else llm.share(prompt=prompt) @NodeConstructor.register('STT') @@ -611,14 +454,12 @@ def cond(x): return True return False - return lazyllm.ifs( - cond, tpath=lazyllm.TrainableModule(base_model), fpath=lazyllm.Identity() - ) + return lazyllm.ifs(cond, tpath=lazyllm.TrainableModule(base_model), fpath=lazyllm.Identity()) @NodeConstructor.register('Constant') def make_constant(value: Any): - return lambda *args, **kw: value + return (lambda *args, **kw: value) class FileResource(object): @@ -631,4 +472,4 @@ def __call__(self, *args, **kw) -> Union[str, List[str]]: @NodeConstructor.register('File') def make_file(id: str): - return FileResource(id) + return FileResource(id) \ No newline at end of file From b1b2e71953ce84f40819166e4e1bc0ae1f64703d Mon Sep 17 00:00:00 2001 From: zhangyongchao Date: Wed, 13 Nov 2024 18:59:40 +0800 Subject: [PATCH 10/12] remove unused import --- tests/basic_tests/test_engine.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/tests/basic_tests/test_engine.py b/tests/basic_tests/test_engine.py index 70c680eb..af4a9732 100644 --- a/tests/basic_tests/test_engine.py +++ b/tests/basic_tests/test_engine.py @@ -1,4 +1,4 @@ -from lazyllm.engine import LightEngine, NodeMetaHook +from lazyllm.engine import LightEngine import pytest import time from gradio_client import Client @@ -7,9 +7,6 @@ from lazyllm.common.common import TimeoutException import json import unittest -from fastapi import FastAPI -from fastapi.responses import JSONResponse -from fastapi.testclient import TestClient import subprocess import socket import threading From 5bb7cdbc344e8aaf0fb5db9d23a4eb8e4d181e9f Mon Sep 17 00:00:00 2001 From: zhangyongchao Date: Wed, 13 Nov 2024 19:02:12 +0800 Subject: [PATCH 11/12] fix lint problem --- lazyllm/engine/engine.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lazyllm/engine/engine.py b/lazyllm/engine/engine.py index 5a0b4553..97ae8d28 100644 --- a/lazyllm/engine/engine.py +++ b/lazyllm/engine/engine.py @@ -472,4 +472,4 @@ def __call__(self, *args, **kw) -> Union[str, List[str]]: @NodeConstructor.register('File') def make_file(id: str): - return FileResource(id) \ No newline at end of file + return FileResource(id) From bbd65e62657d5b4f4daafbfc8839c7068bb680ba Mon Sep 17 00:00:00 2001 From: zhangyongchao Date: Thu, 14 Nov 2024 11:12:53 +0800 Subject: [PATCH 12/12] add get api for data_reflow_hook service --- lazyllm/engine/engine.py | 5 ++- tests/basic_tests/test_engine.py | 55 ++++++++++++++++++++++---------- 2 files changed, 41 insertions(+), 19 deletions(-) diff --git a/lazyllm/engine/engine.py b/lazyllm/engine/engine.py index 97ae8d28..5fbc7c58 100644 --- a/lazyllm/engine/engine.py +++ b/lazyllm/engine/engine.py @@ -134,9 +134,8 @@ def get_args(cls, key, value, builder_key=None): def _process_hook(self, node, module): if not node.enable_data_reflow: return - if not isinstance(module, (lazyllm.ModuleBase, lazyllm.LazyLLMFlowsBase)): - return - node.func.register_hook(NodeMetaHook(node.func, Engine.REPORT_URL, node.id)) + if isinstance(module, (lazyllm.ModuleBase, lazyllm.LazyLLMFlowsBase)): + node.func.register_hook(NodeMetaHook(node.func, Engine.REPORT_URL, node.id)) _constructor = NodeConstructor() diff --git a/tests/basic_tests/test_engine.py b/tests/basic_tests/test_engine.py index af4a9732..6a13893a 100644 --- a/tests/basic_tests/test_engine.py +++ b/tests/basic_tests/test_engine.py @@ -10,23 +10,31 @@ import subprocess import socket import threading +import requests HOOK_PORT = 33733 HOOK_ROUTE = "mock_post" fastapi_code = """ from fastapi import FastAPI from fastapi.responses import JSONResponse -from fastapi.testclient import TestClient +from collections import deque app = FastAPI() +received_datas = deque(maxlen=100) @app.post("/{route}") async def receive_json(data: dict): print("Received json data:", data) - assert "prompt_tokens" in data + received_datas.append(data) return JSONResponse(content=data) +@app.get("/get_last_report") +async def get_last_report(): + if len(received_datas) > 0: + return received_datas[-1] + else: + return {{}} if __name__ == "__main__": import uvicorn @@ -47,6 +55,7 @@ def setUpClass(cls): hostname = socket.gethostname() ip_address = socket.gethostbyname(hostname) cls.report_url = f"http://{ip_address}:{HOOK_PORT}/{HOOK_ROUTE}" + cls.get_url = f"http://{ip_address}:{HOOK_PORT}/get_last_report" def read_stdout(process): for line in iter(process.stdout.readline, b''): @@ -64,6 +73,15 @@ def tearDownClass(cls): cls.fastapi_process.terminate() cls.fastapi_process.wait() + def get_last_report(self): + r = requests.get(self.get_url) + json_obj = {} + try: + json_obj = json.loads(r.content) + except Exception as e: + lazyllm.LOG.warning(str(e)) + return json_obj + @pytest.fixture(autouse=True) def run_around_tests(self): yield @@ -137,6 +155,7 @@ def test_engine_switch(self): assert engine.run(gid, 'case1', 2) == 4 assert engine.run(gid, 'case2', 2) == 6 assert engine.run(gid, 'case3', 3) == 9 + assert "prompt_tokens" in self.get_last_report() def test_engine_ifs(self): plus1 = dict(id='1', kind='Code', name='m1', args=dict(code='def test(x: int):\n return 1 + x\n')) @@ -161,31 +180,32 @@ def test_engine_ifs(self): assert engine.run(gid, 1) == 4 assert engine.run(gid, 5) == 12 assert engine.run(gid, 10) == 100 + assert "prompt_tokens" in self.get_last_report() def test_data_reflow_in_server(self): nodes = [ { - "id": "debug-48f9e95a-e54c-4812-a41b-f61f635816b8-1731324014737", + "id": "1", "kind": "Code", - "name": "1731324014737", + "name": "f1", "args": { "code": "def main(x): return int(x) + 1", "_lazyllm_enable_report": True, }, }, { - "id": "debug-48f9e95a-e54c-4812-a41b-f61f635816b8-1731324018968", + "id": "2", "kind": "Code", - "name": "1731324018968", + "name": "f2", "args": { "code": "def main(x): return int(x) + 2", "_lazyllm_enable_report": True, }, }, { - "id": "debug-48f9e95a-e54c-4812-a41b-f61f635816b8-1731324021877", + "id": "3", "kind": "Code", - "name": "1731324021877", + "name": "f3", "args": { "code": "def main(x): return int(x) + 3", "_lazyllm_enable_report": True, @@ -195,26 +215,26 @@ def test_data_reflow_in_server(self): edges = [ { "iid": "__start__", - "oid": "debug-48f9e95a-e54c-4812-a41b-f61f635816b8-1731324014737", + "oid": "1", }, { - "iid": "debug-48f9e95a-e54c-4812-a41b-f61f635816b8-1731324014737", - "oid": "debug-48f9e95a-e54c-4812-a41b-f61f635816b8-1731324018968", + "iid": "1", + "oid": "2", }, { - "iid": "debug-48f9e95a-e54c-4812-a41b-f61f635816b8-1731324018968", - "oid": "debug-48f9e95a-e54c-4812-a41b-f61f635816b8-1731324021877", + "iid": "2", + "oid": "3", }, { - "iid": "debug-48f9e95a-e54c-4812-a41b-f61f635816b8-1731324021877", + "iid": "3", "oid": "__end__", }, ] resources = [ { - "id": "adc5174c-3220-4b98-92a6-3897e9af8898", + "id": "4", "kind": "server", - "name": "adc5174c-3220-4b98-92a6-3897e9af8898", + "name": "s1", "args": {}, } ] @@ -222,6 +242,7 @@ def test_data_reflow_in_server(self): engine.set_report_url(self.report_url) gid = engine.start(nodes, edges, resources) assert engine.run(gid, 1) == 7 + assert "prompt_tokens" in self.get_last_report() def test_engine_loop(self): nodes = [dict(id='1', kind='Code', name='code', args=dict(code='def square(x: int): return x * x'))] @@ -246,6 +267,7 @@ def test_engine_loop(self): engine.set_report_url(self.report_url) gid = engine.start(nodes, edges) assert engine.run(gid, 2) == 16 + assert "prompt_tokens" in self.get_last_report() def test_engine_warp(self): nodes = [dict(id='1', kind='Code', name='code', args=dict(code='def square(x: int): return x * x'))] @@ -269,6 +291,7 @@ def test_engine_warp(self): engine.set_report_url(self.report_url) gid = engine.start(nodes, edges) assert engine.run(gid, 2, 3, 4, 5) == (4, 9, 16, 25) + assert "prompt_tokens" in self.get_last_report() def test_engine_formatter(self): nodes = [dict(id='1', kind='Formatter', name='f1', args=dict(ftype='python', rule='[:]'))]