add sharedllm and switch tests (#245)

lazyllm/engine/engine.py

@@ -281,3 +281,7 @@ def make_fc(llm: str, tools: List[str], algorithm: Optional[str] = None):
         lazyllm.tools.ReWOOAgent if algorithm == 'ReWOO' else \
         lazyllm.tools.ReactAgent if algorithm == 'React' else lazyllm.tools.FunctionCallAgent
     return f(Engine().build_node(llm).func, tools)
+
+@NodeConstructor.register('SharedLLM')
+def make_shared_llm(llm: str, prompt: Optional[str] = None):
+    return Engine().build_node(llm).func.share(prompt=prompt)

tests/basic_tests/test_engine.py

@@ -12,14 +12,15 @@ def run_around_tests(self):
         LightEngine().reset()
 
     def test_engine_subgraph(self):
-        nodes = [dict(id='1', kind='LocalLLM', name='m1', args=dict(base_model='', deploy_method='dummy'))]
+        resources = [dict(id='0', kind='LocalLLM', name='m1', args=dict(base_model='', deploy_method='dummy'))]
+        nodes = [dict(id='1', kind='SharedLLM', name='s1', args=dict(llm='0', prompt=None))]
         edges = [dict(iid='__start__', oid='1'), dict(iid='1', oid='__end__')]
 
         nodes = [dict(id='2', kind='SubGraph', name='s1', args=dict(nodes=nodes, edges=edges))]
         edges = [dict(iid='__start__', oid='2'), dict(iid='2', oid='__end__')]
 
         engine = LightEngine()
-        engine.start(nodes, edges)
+        engine.start(nodes, edges, resources)
         r = engine.run('1234')
         assert 'reply for You are an AI-Agent developed by LazyLLM' in r
         assert '1234' in r
@@ -50,6 +51,21 @@ def test_engine_switch(self):
         assert engine.run(2) == 6
         assert engine.run(3) == 9
 
+        engine.reset()
+
+        switch = dict(id='4', kind='Switch', name='s1', args=dict(judge_on_full_input=False, nodes={
+            'case1': [double],
+            'case2': [plus1, double],
+            'case3': [square]
+        }))
+        engine.start([switch], edges)
+        assert engine.run('case1', 1) == 2
+        assert engine.run('case2', 1) == 4
+        assert engine.run('case3', 1) == 1
+        assert engine.run('case1', 2) == 4
+        assert engine.run('case2', 2) == 6
+        assert engine.run('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')