From 6c1b56dda064b8475b2f0f2ab4f612b16b1ad5fc Mon Sep 17 00:00:00 2001
From: Artemy Belousov <whimox@gmail.com>
Date: Mon, 2 Sep 2024 16:05:32 +0400
Subject: [PATCH] Add compatibility checks for `config` parameter in
 MotleyTool's LlamaIndex integration & workaround for Pydantic v1
 compatibility (#74)

---
 motleycrew/tools/tool.py | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/motleycrew/tools/tool.py b/motleycrew/tools/tool.py
index be81b5d..7f78d65 100644
--- a/motleycrew/tools/tool.py
+++ b/motleycrew/tools/tool.py
@@ -1,4 +1,5 @@
 import functools
+import inspect
 from typing import Callable, Union, Optional, Dict, Any
 
 from langchain.tools import BaseTool
@@ -157,11 +158,22 @@ def to_llama_index_tool(self) -> LlamaIndex__BaseTool:
             LlamaIndex tool.
         """
         ensure_module_is_installed("llama_index")
+
+        if inspect.signature(self.tool._run).parameters.get("config", None) is not None:
+            fn = functools.partial(self.tool._run, config=RunnableConfig())
+        else:
+            fn = self.tool._run
+
+        fn_schema = self.tool.args_schema
+        fn_schema.model_json_schema = (
+            fn_schema.schema
+        )  # attempt to make it compatible with Langchain's old Pydantic v1
+
         llama_index_tool = LlamaIndex__FunctionTool.from_defaults(
-            fn=functools.partial(self.tool._run, config=RunnableConfig()),
+            fn=fn,
             name=self.tool.name,
             description=self.tool.description,
-            fn_schema=self.tool.args_schema,
+            fn_schema=fn_schema,
         )
         return llama_index_tool