add log_progress

vision_agent/agent/vision_agent.py

@@ -98,6 +98,7 @@ def parse_execution(response: str) -> Optional[str]:
  code = code[: code.find("</execute_python>")]
  return code
 
+
 class VisionAgent(Agent):
  """Vision Agent is an agent that can chat with the user and call tools or other
  agents to generate code for it. Vision Agent uses python code to execute actions for

vision_agent/agent/vision_agent_coder.py

@@ -145,7 +145,8 @@ def pick_plan(
 ) -> Tuple[str, str]:
  log_progress(
  {
- "type": "pick_plan",
+ "type": "log",
+ "log_content": "Generating code to pick the best plan",
  "status": "started",
  }
  )
@@ -162,11 +163,10 @@ def pick_plan(
  code = extract_code(model(prompt))
  log_progress(
  {
- "type": "pick_plan",
+ "type": "log",
+ "log_content": "Executing code to test plans",
+ "code": DefaultImports.prepend_imports(code),
  "status": "running",
- "payload": {
- "code": DefaultImports.prepend_imports(code),
- },
  }
  )
  tool_output = code_interpreter.exec_isolation(DefaultImports.prepend_imports(code))
@@ -178,6 +178,19 @@ def pick_plan(
  _print_code("Initial code and tests:", code)
  _LOGGER.info(f"Initial code execution result:\n{tool_output.text()}")
 
+ log_progress(
+ {
+ "type": "log",
+ "log_content": (
+ "Code execution succeeded"
+ if tool_output.success
+ else "Code execution failed"
+ ),
+ "payload": tool_output.to_json(),
+ "status": "completed" if tool_output.success else "failed",
+ }
+ )
+
  # retry if the tool output is empty or code fails
  count = 0
  while (not tool_output.success or tool_output_str == "") and count < max_retries:
@@ -189,14 +202,26 @@ def pick_plan(
  ),
  media=media,
  )
- code = extract_code(model(prompt))
  log_progress(
  {
- "type": "pick_plan",
+ "type": "log",
+ "log_content": "Retrying code to test plans",
  "status": "running",
- "payload": {
- "code": DefaultImports.prepend_imports(code),
- },
+ "code": DefaultImports.prepend_imports(code),
+ }
+ )
+ code = extract_code(model(prompt))
+ log_progress(
+ {
+ "type": "log",
+ "log_content": (
+ "Code execution succeeded"
+ if tool_output.success
+ else "Code execution failed"
+ ),
+ "code": DefaultImports.prepend_imports(code),
+ "payload": tool_output.to_json(),
+ "status": "completed" if tool_output.success else "failed",
  }
  )
  tool_output = code_interpreter.exec_isolation(
@@ -231,9 +256,10 @@ def pick_plan(
  _LOGGER.info(f"Best plan:\n{best_plan}")
  log_progress(
  {
- "type": "pick_plan",
+ "type": "log",
+ "log_content": "Picked best plan",
  "status": "completed",
- "payload": best_plan,
+ "payload": plans[best_plan["best_plan"]],
  }
  )
  return best_plan["best_plan"], tool_output_str
@@ -304,7 +330,8 @@ def write_and_test_code(
 ) -> Dict[str, Any]:
  log_progress(
  {
- "type": "code",
+ "type": "log",
+ "log_content": "Generating code",
  "status": "started",
  }
  )
@@ -322,10 +349,11 @@ def write_and_test_code(
 
  log_progress(
  {
- "type": "code",
+ "type": "log",
+ "log_content": "Running code",
  "status": "running",
+ "code": DefaultImports.prepend_imports(code),
  "payload": {
- "code": DefaultImports.prepend_imports(code),
  "test": test,
  },
  }
@@ -335,10 +363,15 @@ def write_and_test_code(
  )
  log_progress(
  {
- "type": "code",
+ "type": "log",
+ "log_content": (
+ "Code execution succeeded"
+ if result.success
+ else "Code execution failed"
+ ),
  "status": "completed" if result.success else "failed",
+ "code": DefaultImports.prepend_imports(code),
  "payload": {
- "code": DefaultImports.prepend_imports(code),
  "test": test,
  "result": result.to_json(),
  },
@@ -682,7 +715,8 @@ def chat_with_workflow(
  success = False
  self.log_progress(
  {
- "type": "plans",
+ "type": "log",
+ "log_content": "Creating plans",
  "status": "started",
  }
  )
@@ -698,13 +732,6 @@ def chat_with_workflow(
  _LOGGER.info(
  f"\n{tabulate(tabular_data=plans[p], headers='keys', tablefmt='mixed_grid', maxcolwidths=_MAX_TABULATE_COL_WIDTH)}"
  )
- self.log_progress(
- {
- "type": "plans",
- "status": "completed",
- "payload": plans,
- }
- )
 
  tool_infos = retrieve_tools(
  plans,
@@ -742,11 +769,13 @@ def chat_with_workflow(
 
  self.log_progress(
  {
- "type": "plans",
+ "type": "log",
+ "log_content": "Creating plans",
  "status": "completed",
- "payload": plan_i,
+ "payload": tool_info,
  }
  )
+
  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)}"