fix llama_index output_handler

ShoggothAI · Jun 28, 2024 · bb2a535 · bb2a535
1 parent 092c30c
commit bb2a535
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 11 deletions.
diff --git 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
@@ -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
@@ -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