HttpTool extract result

lazyllm/engine/engine.py

@@ -7,6 +7,7 @@
 from .node_meta_hook import NodeMetaHook
 import inspect
 import functools
+import copy
 from abc import ABC, abstractclassmethod
 
 
@@ -16,8 +17,9 @@ class Engine(ABC):
     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: Dict[str, Node] = {
+            '__start__': Node(id='__start__', kind='__start__', name='__start__'),
+            '__end__': Node(id='__end__', kind='__end__', name='__end__')}
 
     def __new__(cls):
         if cls is not Engine:
@@ -103,13 +105,14 @@ 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_args = copy.copy(node.args)
+        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
 
@@ -122,7 +125,7 @@ def get_args(cls, key, value, builder_key=None):
                 return Engine().build_node(value).func
             return node_args.getattr_f(value) if node_args.getattr_f else value
 
-        for key, value in node.args.items():
+        for key, value in node_args.items():
             if key in node_msgs['init_arguments']:
                 init_args[key] = get_args('init_arguments', key, value)
             elif key in node_msgs['builder_argument']:
@@ -418,9 +421,11 @@ def make_http_tool(method: Optional[str] = None,
                    proxies: Optional[Dict[str, str]] = None,
                    code_str: Optional[str] = None,
                    vars_for_code: Optional[Dict[str, Any]] = None,
-                   doc: Optional[str] = None):
+                   doc: Optional[str] = None,
+                   outputs: Optional[List[str]] = None,
+                   extract_from_result: Optional[bool] = None):
     instance = lazyllm.tools.HttpTool(method, url, params, headers, body, timeout, proxies,
-                                      code_str, vars_for_code)
+                                      code_str, vars_for_code, outputs, extract_from_result)
     if doc:
         instance.__doc__ = doc
     return instance

lazyllm/tools/tools/http_tool.py

@@ -1,6 +1,7 @@
-from lazyllm.common import compile_func
+from lazyllm.common import compile_func, package
 from lazyllm.tools.http_request import HttpRequest
-from typing import Optional, Dict, Any
+from typing import Optional, Dict, Any, List
+import json
 
 class HttpTool(HttpRequest):
     def __init__(self,
@@ -12,15 +13,35 @@ def __init__(self,
                  timeout: int = 10,
                  proxies: Optional[Dict[str, str]] = None,
                  code_str: Optional[str] = None,
-                 vars_for_code: Optional[Dict[str, Any]] = None):
+                 vars_for_code: Optional[Dict[str, Any]] = None,
+                 outputs: Optional[List[str]] = None,
+                 extract_from_result: Optional[bool] = None):
         super().__init__(method, url, '', headers, params, body, timeout, proxies)
         self._has_http = True if url else False
-        self._compiled_func = compile_func(code_str, vars_for_code) if code_str else None
+        self._compiled_func = (compile_func(code_str, vars_for_code) if code_str else
+                               (lambda x: json.loads(x['content'])) if self._has_http else None)
+        self._outputs, self._extract_from_result = outputs, extract_from_result
+        if extract_from_result:
+            assert outputs, 'Output information is necessary to extract output parameters'
+            assert len(outputs) == 1, 'When the number of outputs is greater than 1, no manual setting is required'
+
+    def _get_result(self, res):
+        if self._extract_from_result or (isinstance(res, dict) and len(self._outputs) > 1):
+            assert isinstance(res, dict), 'The result of the tool should be a dict type'
+            r = package(res.get(key) for key in self._outputs)
+            return r[0] if len(r) == 1 else r
+        if len(self._outputs) > 1:
+            assert isinstance(res, (tuple, list)), 'The result of the tool should be tuple or list'
+            assert len(res) == len(self._outputs), 'The number of outputs is inconsistent with expectations'
+            return package(res)
+        return res
 
     def forward(self, *args, **kwargs):
+        if not self._compiled_func: return None
         if self._has_http:
             res = super().forward(*args, **kwargs)
-            return self._compiled_func(res) if self._compiled_func else res
-        elif self._compiled_func:
-            return self._compiled_func(*args, **kwargs)
-        return None
+            if int(res['status_code']) >= 400:
+                raise RuntimeError(f'HttpRequest error, status code is {res["status_code"]}.')
+            args, kwargs = (res,), {}
+        res = self._compiled_func(*args, **kwargs)
+        return self._get_result(res) if self._outputs else res

tests/basic_tests/test_engine.py

@@ -567,10 +567,48 @@ def test_engine_httptool(self):
         engine = LightEngine()
         gid = engine.start(nodes, edges, gid='graph-1')
         res = engine.run(gid)
-        content = json.loads(res['content'])
 
-        assert content['headers']['h1'] == 'baz'
-        assert content['url'].endswith(f'{url[5:]}?p1=foo&p2=bar')
+        assert res['headers']['h1'] == 'baz'
+        assert res['url'].endswith(f'{url[5:]}?p1=foo&p2=bar')
+
+    def test_engine_httptool_with_output(self):
+        params = {'p1': '{{p1}}', 'p2': '{{p2}}'}
+        headers = {'h1': '{{h1}}'}
+        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,
+                outputs=['headers', 'url'], _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')
+        res = engine.run(gid)
+
+        assert isinstance(res, lazyllm.package) and len(res) == 2
+        assert res[0]['h1'] == 'baz'
+        assert res[1].endswith(f'{url[5:]}?p1=foo&p2=bar')
+
+        engine.reset()
+
+        nodes[3]['args']['outputs'] = ['output']
+        gid = engine.start(nodes, edges)
+        res = engine.run(gid)
+        assert res['headers']['h1'] == 'baz'
+
+        engine.reset()
+
+        nodes[3]['args']['outputs'] = ['headers']
+        nodes[3]['args']['extract_from_result'] = True
+        gid2 = engine.start(nodes, edges)
+        res = engine.run(gid2)
+        assert res['h1'] == 'baz'
 
     def test_engine_httptool_body(self):
         body = {'b1': '{{b1}}', 'b2': '{{b2}}'}
@@ -590,10 +628,9 @@ def test_engine_httptool_body(self):
         engine = LightEngine()
         gid = engine.start(nodes, edges, gid='graph-1')
         res = engine.run(gid)
-        content = json.loads(res['content'])
 
-        assert content['b1'] == 'body1'
-        assert content['b2'] == 'body2'
+        assert res['b1'] == 'body1'
+        assert res['b2'] == 'body2'
 
     def test_engine_status(self):
         resources = [dict(id='0', kind='LocalLLM', name='m1', args=dict(base_model='', deploy_method='dummy'))]