From bb2a535be405e0e48417ba45153d86e0724f40c6 Mon Sep 17 00:00:00 2001
From: User <serfintw@yadnex.ru>
Date: Fri, 28 Jun 2024 16:41:09 +0300
Subject: [PATCH] fix llama_index output_handler

---
 motleycrew/agents/llama_index/llama_index.py |  2 ++
 motleycrew/agents/parent.py                  | 19 +++++++++++--------
 motleycrew/tracking/callbacks.py             | 15 ++++++++++++---
 3 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/motleycrew/agents/llama_index/llama_index.py b/motleycrew/agents/llama_index/llama_index.py
index 47dc5700..5ae9fd0d 100644
--- a/motleycrew/agents/llama_index/llama_index.py
+++ b/motleycrew/agents/llama_index/llama_index.py
@@ -24,6 +24,8 @@
 
 
 class LlamaIndexMotleyAgent(MotleyAgentParent):
+    wrapp_direct_tool = False
+
     def __init__(
         self,
         description: str | None = None,
diff --git a/motleycrew/agents/parent.py b/motleycrew/agents/parent.py
index f0a6d34c..0b2cf0de 100644
--- a/motleycrew/agents/parent.py
+++ b/motleycrew/agents/parent.py
@@ -55,6 +55,8 @@ def wrapper(*args, **kwargs):
 
 
 class MotleyAgentParent(MotleyAgentAbstractParent, Runnable):
+    wrapp_direct_tool = True
+
     def __init__(
         self,
         description: str | None = None,
@@ -175,15 +177,16 @@ def handle_agent_output(*args, **kwargs):
             args_schema=self.output_handler.args_schema,
         )
 
-        object.__setattr__(
-            prepared_output_handler,
-            "_run",
-            _run_tool_direct_decorator(prepared_output_handler._run),
-        )
+        if self.wrapp_direct_tool:
+            object.__setattr__(
+                prepared_output_handler,
+                "_run",
+                _run_tool_direct_decorator(prepared_output_handler._run),
+            )
 
-        object.__setattr__(
-            prepared_output_handler, "run", run_tool_direct_decorator(prepared_output_handler.run)
-        )
+            object.__setattr__(
+                prepared_output_handler, "run", run_tool_direct_decorator(prepared_output_handler.run)
+            )
 
         return MotleyTool.from_langchain_tool(prepared_output_handler)
 
diff --git a/motleycrew/tracking/callbacks.py b/motleycrew/tracking/callbacks.py
index f03d5b2e..8e4c7317 100644
--- a/motleycrew/tracking/callbacks.py
+++ b/motleycrew/tracking/callbacks.py
@@ -231,7 +231,11 @@ def _on_function_call_start(
 
         params["parent_run_id"] = self.check_parent_id(parent_id)
         params["name"] = tool.name
-        params["input"] = '{{"query":"{}"}}'.format(function_call.get("query"))
+
+        params_inputs = []
+        for k, v in function_call.items():
+            params_inputs.append('{{"{}":"{}"}}'.format(k, v))
+        params["input"] = "\n".join(params_inputs)
 
         return params
 
@@ -302,11 +306,16 @@ def _on_agent_step_end(
         Returns:
             dict: dictionary of data to send to lunary
         """
-        response = payload.get(EventPayload.RESPONSE)
         params = self._get_initial_track_event_params(
             LunaryRunType.AGENT, LunaryEventName.END, event_id
         )
-        params["output"] = response.response
+
+        if payload:
+            response = payload.get(EventPayload.RESPONSE)
+            output  = response.response
+        else:
+            output = ""
+        params["output"] = output
 
         return params