fix(query_exec_tracker): tracker not adding the code correctly (#804)

sinaptik-ai · Dec 10, 2023 · 250c5fe · 250c5fe
1 parent 26857ef
commit 250c5fe
Showing 1 changed file with 2 additions and 5 deletions.
diff --git a/pandasai/helpers/query_exec_tracker.py b/pandasai/helpers/query_exec_tracker.py
@@ -174,18 +174,15 @@ def _generate_exec_step(self, func_name: str, result: Any) -> dict:
             step[
                 "type"
             ] = f"{exec_steps[func_name]} ({self._func_exec_count['_retry_run_code']})"
-            step["code_generated"] = result[0]
+            step["code_generated"] = result
 
-        elif func_name == "cache_hit":
+        elif func_name in {"cache_hit", "generate_code"}:
             step["code_generated"] = result
 
         elif func_name == "execute_code":
             self._response = self._format_response(result)
             step["result"] = self._response
 
-        elif func_name == "generate_code":
-            step["code_generated"] = result[0]
-
         return step
 
     def _format_response(self, result: ResponseType) -> ResponseType: