Skip to content

Commit

Permalink
fix bug: map backend_nodeid to frontend_nodeid (#341)
Browse files Browse the repository at this point in the history
  • Loading branch information
SuperEver authored Nov 12, 2024
1 parent 96640c9 commit f8e1510
Show file tree
Hide file tree
Showing 4 changed files with 63 additions and 60 deletions.
11 changes: 8 additions & 3 deletions lazyllm/engine/engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
3 changes: 1 addition & 2 deletions lazyllm/engine/node.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,15 +9,14 @@

@dataclass
class Node():
id: int
id: str
kind: str
name: str
args: Optional[Dict] = None

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
Expand Down
102 changes: 54 additions & 48 deletions tests/basic_tests/test_engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +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='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()
Expand All @@ -69,9 +69,10 @@ 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=dict(code='def test(x: int):\n return x * x\n', _lazyllm_enable_report=True))
switch = dict(
id="4",
kind="Switch",
Expand Down Expand Up @@ -112,9 +113,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')
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",
Expand All @@ -136,7 +137,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='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 = [
Expand All @@ -160,7 +161,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='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 = [
Expand Down Expand Up @@ -222,10 +223,12 @@ def test_engine_formatter(self):
'<lazyllm-query>{"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__')]
Expand All @@ -236,9 +239,9 @@ 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__')]

Expand All @@ -248,11 +251,13 @@ 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'),
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='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')]
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__')]

Expand All @@ -264,8 +269,9 @@ 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'),
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]')]
Expand All @@ -276,7 +282,7 @@ 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'),
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()
Expand All @@ -287,9 +293,9 @@ 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'),
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__')]
Expand All @@ -299,7 +305,7 @@ 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]'),
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()
Expand All @@ -310,9 +316,9 @@ 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'),
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__')]
Expand All @@ -322,9 +328,9 @@ 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'),
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__')]
Expand All @@ -335,9 +341,9 @@ 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]
))
Expand All @@ -349,7 +355,7 @@ 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]
))
Expand All @@ -361,9 +367,9 @@ 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'),
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__')]
Expand All @@ -376,7 +382,7 @@ 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))
Expand Down Expand Up @@ -433,9 +439,9 @@ 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='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']))
]
Expand All @@ -454,9 +460,9 @@ 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(
Expand Down Expand Up @@ -496,7 +502,7 @@ def test_rag(self):
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}\''),
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'),
Expand Down
7 changes: 0 additions & 7 deletions tests/charge_tests/test_sql_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()

0 comments on commit f8e1510

Please sign in to comment.