12
12
# See the License for the specific language governing permissions and
13
13
# limitations under the License.
14
14
15
- import asyncio
15
+ from asyncio import wrap_future
16
16
from typing import Any , Callable , Union
17
17
18
18
from langchain_core .tools import BaseTool
19
19
from toolbox_core .sync_tool import ToolboxSyncTool as ToolboxCoreSyncTool
20
+ from toolbox_core .utils import params_to_pydantic_model
20
21
21
22
22
23
class ToolboxTool (BaseTool ):
@@ -41,24 +42,15 @@ def __init__(
41
42
super ().__init__ (
42
43
name = core_tool .__name__ ,
43
44
description = core_tool .__doc__ ,
44
- args_schema = core_tool ._async_tool . _pydantic_model ,
45
+ args_schema = params_to_pydantic_model ( core_tool ._name , core_tool . _params ) ,
45
46
)
46
47
self .__core_tool = core_tool
47
48
48
49
def _run (self , ** kwargs : Any ) -> str :
49
50
return self .__core_tool (** kwargs )
50
51
51
52
async def _arun (self , ** kwargs : Any ) -> str :
52
- coro = self .__core_tool ._async_tool (** kwargs )
53
-
54
- # If a loop has not been provided, attempt to run in current thread.
55
- if not self .__core_tool ._loop :
56
- return await coro
57
-
58
- # Otherwise, run in the background thread.
59
- return await asyncio .wrap_future (
60
- asyncio .run_coroutine_threadsafe (coro , self .__core_tool ._loop )
61
- )
53
+ return await wrap_future (self .__core_tool ._call_future (** kwargs ))
62
54
63
55
def add_auth_token_getters (
64
56
self , auth_token_getters : dict [str , Callable [[], str ]]
0 commit comments