From c8fc196c3a952e1a2353c64b63da4dc01bf31ede Mon Sep 17 00:00:00 2001
From: wuyiqunLu <132986242+wuyiqunLu@users.noreply.github.com>
Date: Mon, 22 Jul 2024 14:02:07 +0800
Subject: [PATCH] fix: add back some logging for ui (#176)

* fix: add back some logging for ui

* simplifier

* resolve lint error

* fix

* fix type error

* fix type error

* fix type error

* remove check
---
 vision_agent/agent/vision_agent.py | 89 ++++++++++++++++--------------
 1 file changed, 49 insertions(+), 40 deletions(-)

diff --git a/vision_agent/agent/vision_agent.py b/vision_agent/agent/vision_agent.py
index 110a9a17..d8fc079a 100644
--- a/vision_agent/agent/vision_agent.py
+++ b/vision_agent/agent/vision_agent.py
@@ -172,19 +172,25 @@ def write_plans(
 def pick_plan(
     chat: List[Message],
     plans: Dict[str, Any],
-    tool_info: str,
+    tool_infos: Dict[str, str],
     model: LMM,
     code_interpreter: CodeInterpreter,
+    test_multi_plan: bool,
     verbosity: int = 0,
     max_retries: int = 3,
-) -> Tuple[str, str]:
+) -> Tuple[Any, str, str]:
+    if not test_multi_plan:
+        k = list(plans.keys())[0]
+        return plans[k], tool_infos[k], ""
+
+    all_tool_info = tool_infos["all"]
     chat = copy.deepcopy(chat)
     if chat[-1]["role"] != "user":
         raise ValueError("Last chat message must be from the user.")
 
     plan_str = format_plans(plans)
     prompt = TEST_PLANS.format(
-        docstring=tool_info, plans=plan_str, previous_attempts=""
+        docstring=all_tool_info, plans=plan_str, previous_attempts=""
     )
 
     code = extract_code(model(prompt))
@@ -201,7 +207,7 @@ def pick_plan(
     count = 0
     while (not tool_output.success or tool_output_str == "") and count < max_retries:
         prompt = TEST_PLANS.format(
-            docstring=tool_info,
+            docstring=all_tool_info,
             plans=plan_str,
             previous_attempts=PREVIOUS_FAILED.format(
                 code=code, error=tool_output.text()
@@ -237,7 +243,17 @@ def pick_plan(
     best_plan = extract_json(model(chat))
     if verbosity >= 1:
         _LOGGER.info(f"Best plan:\n{best_plan}")
-    return best_plan["best_plan"], tool_output_str
+
+    plan = best_plan["best_plan"]
+    if plan in plans and plan in tool_infos:
+        return plans[plan], tool_infos[plan], tool_output_str
+    else:
+        if verbosity >= 1:
+            _LOGGER.warning(
+                f"Best plan {plan} not found in plans or tool_infos. Using the first plan and tool info."
+            )
+        k = list(plans.keys())[0]
+        return plans[k], tool_infos[k], tool_output_str
 
 
 @traceable
@@ -524,6 +540,13 @@ def retrieve_tools(
         )
     all_tools = "\n\n".join(set(tool_info))
     tool_lists_unique["all"] = all_tools
+    log_progress(
+        {
+            "type": "tools",
+            "status": "completed",
+            "payload": tool_lists[list(plans.keys())[0]],
+        }
+    )
     return tool_lists_unique
 
 
@@ -692,6 +715,14 @@ def chat_with_workflow(
                 self.planner,
             )
 
+            self.log_progress(
+                {
+                    "type": "plans",
+                    "status": "completed",
+                    "payload": plans[list(plans.keys())[0]],
+                }
+            )
+
             if self.verbosity >= 1 and test_multi_plan:
                 for p in plans:
                     _LOGGER.info(
@@ -705,47 +736,25 @@ def chat_with_workflow(
                 self.verbosity,
             )
 
-            if test_multi_plan:
-                best_plan, tool_output_str = pick_plan(
-                    int_chat,
-                    plans,
-                    tool_infos["all"],
-                    self.coder,
-                    code_interpreter,
-                    verbosity=self.verbosity,
-                )
-            else:
-                best_plan = list(plans.keys())[0]
-                tool_output_str = ""
-
-            if best_plan in plans and best_plan in tool_infos:
-                plan_i = plans[best_plan]
-                tool_info = tool_infos[best_plan]
-            else:
-                if self.verbosity >= 1:
-                    _LOGGER.warning(
-                        f"Best plan {best_plan} not found in plans or tool_infos. Using the first plan and tool info."
-                    )
-                k = list(plans.keys())[0]
-                plan_i = plans[k]
-                tool_info = tool_infos[k]
-
-            self.log_progress(
-                {
-                    "type": "plans",
-                    "status": "completed",
-                    "payload": plan_i,
-                }
+            best_plan, best_tool_info, tool_output_str = pick_plan(
+                int_chat,
+                plans,
+                tool_infos,
+                self.coder,
+                code_interpreter,
+                test_multi_plan,
+                verbosity=self.verbosity,
             )
+
             if self.verbosity >= 1:
                 _LOGGER.info(
-                    f"Picked best plan:\n{tabulate(tabular_data=plan_i, headers='keys', tablefmt='mixed_grid', maxcolwidths=_MAX_TABULATE_COL_WIDTH)}"
+                    f"Picked best plan:\n{tabulate(tabular_data=best_plan, headers='keys', tablefmt='mixed_grid', maxcolwidths=_MAX_TABULATE_COL_WIDTH)}"
                 )
 
             results = write_and_test_code(
                 chat=[{"role": c["role"], "content": c["content"]} for c in int_chat],
-                plan="\n-" + "\n-".join([e["instructions"] for e in plan_i]),
-                tool_info=tool_info,
+                plan="\n-" + "\n-".join([e["instructions"] for e in best_plan]),
+                tool_info=best_tool_info,
                 tool_output=tool_output_str,
                 tool_utils=T.UTILITIES_DOCSTRING,
                 working_memory=working_memory,
@@ -761,7 +770,7 @@ def chat_with_workflow(
             code = cast(str, results["code"])
             test = cast(str, results["test"])
             working_memory.extend(results["working_memory"])  # type: ignore
-            plan.append({"code": code, "test": test, "plan": plan_i})
+            plan.append({"code": code, "test": test, "plan": best_plan})
 
             execution_result = cast(Execution, results["test_result"])
             self.log_progress(