vision-agent 0.2.30__py3-none-any.whl → 0.2.32__py3-none-any.whl

vision_agent/agent/vision_agent.py

@@ -1,778 +1,456 @@
+import copy
 import json
 import logging
 import sys
-import tempfile
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Union, cast
 
-from PIL import Image
+from rich.console import Console
+from rich.syntax import Syntax
 from tabulate import tabulate
 
-from vision_agent.agent.agent import Agent
-from vision_agent.agent.easytool_prompts import (
-    ANSWER_GENERATE,
-    ANSWER_SUMMARIZE,
-    CHOOSE_PARAMETER,
-    CHOOSE_TOOL,
-    TASK_DECOMPOSE,
-    TASK_TOPOLOGY,
-)
+from vision_agent.agent import Agent
 from vision_agent.agent.vision_agent_prompts import (
-    ANSWER_GENERATE_DEPENDS,
-    ANSWER_SUMMARIZE_DEPENDS,
-    CHOOSE_PARAMETER_DEPENDS,
-    CHOOSE_TOOL_DEPENDS,
-    TASK_DECOMPOSE_DEPENDS,
-    VISION_AGENT_REFLECTION,
+    CODE,
+    FEEDBACK,
+    FIX_BUG,
+    FULL_TASK,
+    PLAN,
+    REFLECT,
+    SIMPLE_TEST,
+    USER_REQ,
 )
 from vision_agent.llm import LLM, OpenAILLM
 from vision_agent.lmm import LMM, OpenAILMM
-from vision_agent.tools import TOOLS
-from vision_agent.utils.image_utils import (
-    convert_to_b64,
-    overlay_bboxes,
-    overlay_heat_map,
-    overlay_masks,
-)
+from vision_agent.tools import TOOL_DESCRIPTIONS, TOOLS_DF, UTILITIES_DOCSTRING
+from vision_agent.utils import Execute
+from vision_agent.utils.sim import Sim
 
 logging.basicConfig(stream=sys.stdout)
 _LOGGER = logging.getLogger(__name__)
 _MAX_TABULATE_COL_WIDTH = 80
+_EXECUTE = Execute(600)
+_CONSOLE = Console()
 
 
-def parse_json(s: str) -> Any:
-    s = (
-        s.replace(": True", ": true")
-        .replace(": False", ": false")
-        .replace(":True", ": true")
-        .replace(":False", ": false")
-        .replace("```", "")
-        .strip()
-    )
-    return json.loads(s)
-
-
-def change_name(name: str) -> str:
-    change_list = ["from", "class", "return", "false", "true", "id", "and", "", "ID"]
-    if name in change_list:
-        name = "is_" + name.lower()
-    return name
-
-
-def format_tools(tools: Dict[int, Any]) -> str:
-    # Format this way so it's clear what the ID's are
-    tool_str = ""
-    for key in tools:
-        tool_str += f"ID: {key} - {tools[key]}\n"
-    return tool_str
-
-
-def format_tool_usage(tools: Dict[int, Any], tool_result: List[Dict]) -> str:
-    usage = []
-    name_to_usage = {v["name"]: v["usage"] for v in tools.values()}
-    for tool_res in tool_result:
-        if "tool_name" in tool_res:
-            usage.append((tool_res["tool_name"], name_to_usage[tool_res["tool_name"]]))
-
-    usage_str = ""
-    for tool_name, tool_usage in usage:
-        usage_str += f"{tool_name} - {tool_usage}\n"
-    return usage_str
-
-
-def topological_sort(tasks: List[Dict]) -> List[Dict]:
-    in_degree = {task["id"]: 0 for task in tasks}
-    for task in tasks:
-        for dep in task["dep"]:
-            if dep in in_degree:
-                in_degree[task["id"]] += 1
-
-    queue = [task for task in tasks if in_degree[task["id"]] == 0]
-    sorted_order = []
-
-    while queue:
-        current = queue.pop(0)
-        sorted_order.append(current)
-
-        for task in tasks:
-            if current["id"] in task["dep"]:
-                in_degree[task["id"]] -= 1
-                if in_degree[task["id"]] == 0:
-                    queue.append(task)
-
-    if len(sorted_order) != len(tasks):
-        completed_ids = set([task["id"] for task in sorted_order])
-        remaining_tasks = [task for task in tasks if task["id"] not in completed_ids]
-        sorted_order.extend(remaining_tasks)
-    return sorted_order
-
-
-def task_decompose(
-    model: Union[LLM, LMM, Agent],
-    question: str,
-    tools: Dict[int, Any],
-    reflections: str,
-) -> Optional[Dict]:
-    if reflections:
-        prompt = TASK_DECOMPOSE_DEPENDS.format(
-            question=question, tools=format_tools(tools), reflections=reflections
-        )
-    else:
-        prompt = TASK_DECOMPOSE.format(question=question, tools=format_tools(tools))
-    tries = 0
-    str_result = ""
-    while True:
-        try:
-            str_result = model(prompt)
-            result = parse_json(str_result)
-            return result["Tasks"]  # type: ignore
-        except Exception:
-            if tries > 10:
-                _LOGGER.error(f"Failed task_decompose on: {str_result}")
-                return None
-            tries += 1
-            continue
-
-
-def task_topology(
-    model: Union[LLM, LMM, Agent], question: str, task_list: List[Dict]
-) -> List[Dict[str, Any]]:
-    prompt = TASK_TOPOLOGY.format(question=question, task_list=task_list)
-    tries = 0
-    str_result = ""
-    while True:
-        try:
-            str_result = model(prompt)
-            result = parse_json(str_result)
-            for elt in result["Tasks"]:
-                if isinstance(elt["dep"], str):
-                    elt["dep"] = [int(dep) for dep in elt["dep"].split(",")]
-                elif isinstance(elt["dep"], int):
-                    elt["dep"] = [elt["dep"]]
-                elif isinstance(elt["dep"], list):
-                    elt["dep"] = [int(dep) for dep in elt["dep"]]
-            return result["Tasks"]  # type: ignore
-        except Exception:
-            if tries > 10:
-                _LOGGER.error(f"Failed task_topology on: {str_result}")
-                return task_list
-            tries += 1
-            continue
-
-
-def choose_tool(
-    model: Union[LLM, LMM, Agent],
-    question: str,
-    tools: Dict[int, Any],
-    reflections: str,
-) -> Optional[int]:
-    if reflections:
-        prompt = CHOOSE_TOOL_DEPENDS.format(
-            question=question, tools=format_tools(tools), reflections=reflections
-        )
-    else:
-        prompt = CHOOSE_TOOL.format(question=question, tools=format_tools(tools))
-    tries = 0
-    str_result = ""
-    while True:
-        try:
-            str_result = model(prompt)
-            result = parse_json(str_result)
-            return result["ID"]  # type: ignore
-        except Exception:
-            if tries > 10:
-                _LOGGER.error(f"Failed choose_tool on: {str_result}")
-                return None
-            tries += 1
-            continue
-
-
-def choose_parameter(
-    model: Union[LLM, LMM, Agent],
-    question: str,
-    tool_usage: Dict,
-    previous_log: str,
-    reflections: str,
-) -> Optional[Any]:
-    # TODO: should format tool_usage
-    if reflections:
-        prompt = CHOOSE_PARAMETER_DEPENDS.format(
-            question=question,
-            tool_usage=tool_usage,
-            previous_log=previous_log,
-            reflections=reflections,
-        )
-    else:
-        prompt = CHOOSE_PARAMETER.format(
-            question=question, tool_usage=tool_usage, previous_log=previous_log
-        )
-    tries = 0
-    str_result = ""
-    while True:
-        try:
-            str_result = model(prompt)
-            result = parse_json(str_result)
-            return result["Parameters"]
-        except Exception:
-            if tries > 10:
-                _LOGGER.error(f"Failed choose_parameter on: {str_result}")
-                return None
-            tries += 1
-            continue
-
-
-def answer_generate(
-    model: Union[LLM, LMM, Agent],
-    question: str,
-    call_results: str,
-    previous_log: str,
-    reflections: str,
-) -> str:
-    if reflections:
-        prompt = ANSWER_GENERATE_DEPENDS.format(
-            question=question,
-            call_results=call_results,
-            previous_log=previous_log,
-            reflections=reflections,
-        )
-    else:
-        prompt = ANSWER_GENERATE.format(
-            question=question, call_results=call_results, previous_log=previous_log
+def format_memory(memory: List[Dict[str, str]]) -> str:
+    return FEEDBACK.format(
+        feedback="\n".join(
+            [
+                f"### Feedback {i}:\nCode: ```python\n{m['code']}\n```\nFeedback: {m['feedback']}\n"
+                for i, m in enumerate(memory)
+            ]
         )
-    return model(prompt)
+    )
 
 
-def answer_summarize(
-    model: Union[LLM, LMM, Agent], question: str, answers: List[Dict], reflections: str
-) -> str:
-    if reflections:
-        prompt = ANSWER_SUMMARIZE_DEPENDS.format(
-            question=question, answers=answers, reflections=reflections
-        )
+def extract_code(code: str) -> str:
+    if "\n```python" in code:
+        start = "\n```python"
+    elif "```python" in code:
+        start = "```python"
     else:
-        prompt = ANSWER_SUMMARIZE.format(question=question, answers=answers)
-    return model(prompt)
+        return code
+
+    code = code[code.find(start) + len(start) :]
+    code = code[: code.find("```")]
+    if code.startswith("python\n"):
+        code = code[len("python\n") :]
+    return code
 
 
-def function_call(tool: Callable, parameters: Dict[str, Any]) -> Any:
+def extract_json(json_str: str) -> Dict[str, Any]:
     try:
-        return tool()(**parameters)
-    except Exception as e:
-        _LOGGER.error(f"Failed function_call on: {e}")
-        # return error message so it can self-correct
-        return str(e)
-
-
-def self_reflect(
-    reflect_model: Union[LLM, LMM],
-    question: str,
-    tools: Dict[int, Any],
-    tool_result: List[Dict],
-    final_answer: str,
-    images: Optional[Sequence[Union[str, Path]]] = None,
-) -> str:
-    prompt = VISION_AGENT_REFLECTION.format(
-        question=question,
-        tools=format_tools({k: v["description"] for k, v in tools.items()}),
-        tool_usage=format_tool_usage(tools, tool_result),
-        tool_results=str(tool_result),
-        final_answer=final_answer,
+        json_dict = json.loads(json_str)
+    except json.JSONDecodeError:
+        if "```json" in json_str:
+            json_str = json_str[json_str.find("```json") + len("```json") :]
+            json_str = json_str[: json_str.find("```")]
+        elif "```" in json_str:
+            json_str = json_str[json_str.find("```") + len("```") :]
+            # get the last ``` not one from an intermediate string
+            json_str = json_str[: json_str.find("}```")]
+        json_dict = json.loads(json_str)
+    return json_dict  # type: ignore
+
+
+def write_plan(
+    chat: List[Dict[str, str]],
+    tool_desc: str,
+    working_memory: str,
+    model: Union[LLM, LMM],
+    media: Optional[List[Union[str, Path]]] = None,
+) -> List[Dict[str, str]]:
+    chat = copy.deepcopy(chat)
+    if chat[-1]["role"] != "user":
+        raise ValueError("Last chat message must be from the user.")
+
+    user_request = chat[-1]["content"]
+    context = USER_REQ.format(user_request=user_request)
+    prompt = PLAN.format(context=context, tool_desc=tool_desc, feedback=working_memory)
+    chat[-1]["content"] = prompt
+    if isinstance(model, OpenAILMM):
+        return extract_json(model.chat(chat, images=media))["plan"]  # type: ignore
+    else:
+        return extract_json(model.chat(chat))["plan"]  # type: ignore
+
+
+def reflect(
+    chat: List[Dict[str, str]],
+    plan: str,
+    code: str,
+    model: LLM,
+) -> Dict[str, Union[str, bool]]:
+    chat = copy.deepcopy(chat)
+    if chat[-1]["role"] != "user":
+        raise ValueError("Last chat message must be from the user.")
+
+    user_request = chat[-1]["content"]
+    context = USER_REQ.format(user_request=user_request)
+    prompt = REFLECT.format(context=context, plan=plan, code=code)
+    chat[-1]["content"] = prompt
+    return extract_json(model.chat(chat))
+
+
+def write_and_test_code(
+    task: str,
+    tool_info: str,
+    tool_utils: str,
+    working_memory: str,
+    coder: LLM,
+    tester: LLM,
+    debugger: LLM,
+    log_progress: Callable[[Dict[str, Any]], None],
+    verbosity: int = 0,
+    max_retries: int = 3,
+    input_media: Optional[Union[str, Path]] = None,
+) -> Dict[str, Any]:
+    code = extract_code(
+        coder(CODE.format(docstring=tool_info, question=task, feedback=working_memory))
+    )
+    test = extract_code(
+        tester(
+            SIMPLE_TEST.format(
+                docstring=tool_utils,
+                question=task,
+                code=code,
+                feedback=working_memory,
+                media=input_media,
+            )
+        )
     )
-    if (
-        issubclass(type(reflect_model), LMM)
-        and images is not None
-        and all([Path(image).suffix in [".jpg", ".jpeg", ".png"] for image in images])
-    ):
-        return reflect_model(prompt, images=images)  # type: ignore
-    return reflect_model(prompt)
 
+    success, result = _EXECUTE.run_isolation(f"{code}\n{test}")
+    if verbosity == 2:
+        _LOGGER.info("Initial code and tests:")
+        log_progress(
+            {
+                "log": "Code:",
+                "code": code,
+            }
+        )
+        log_progress(
+            {
+                "log": "Test:",
+                "code": test,
+            }
+        )
+        _CONSOLE.print(
+            Syntax(f"{code}\n{test}", "python", theme="gruvbox-dark", line_numbers=True)
+        )
+        log_progress(
+            {
+                "log": "Result:",
+                "result": result,
+            }
+        )
+        _LOGGER.info(f"Initial result: {result}")
+
+    count = 0
+    new_working_memory = []
+    while not success and count < max_retries:
+        fixed_code_and_test = extract_json(
+            debugger(
+                FIX_BUG.format(
+                    code=code, tests=test, result=result, feedback=working_memory
+                )
+            )
+        )
+        if fixed_code_and_test["code"].strip() != "":
+            code = extract_code(fixed_code_and_test["code"])
+        if fixed_code_and_test["test"].strip() != "":
+            test = extract_code(fixed_code_and_test["test"])
+        new_working_memory.append(
+            {"code": f"{code}\n{test}", "feedback": fixed_code_and_test["reflections"]}
+        )
 
-def parse_reflect(reflect: str) -> Any:
-    reflect = reflect.strip()
-    try:
-        return parse_json(reflect)
-    except Exception:
-        _LOGGER.error(f"Failed parse json reflection: {reflect}")
-    # LMMs have a hard time following directions, so make the criteria less strict
-    finish = (
-        "finish" in reflect.lower() and len(reflect) < 100
-    ) or "finish" in reflect.lower()[-10:]
-    return {"Finish": finish, "Reflection": reflect}
-
-
-def _handle_extract_frames(
-    image_to_data: Dict[str, Dict], tool_result: Dict
-) -> Dict[str, Dict]:
-    image_to_data = image_to_data.copy()
-    # handle extract_frames_ case, useful if it extracts frames but doesn't do
-    # any following processing
-    for video_file_output in tool_result["call_results"]:
-        # When the video tool is run with wrong parameters, exit the loop
-        if not isinstance(video_file_output, tuple) or len(video_file_output) < 2:
-            break
-        for frame, _ in video_file_output:
-            image = frame
-            if image not in image_to_data:
-                image_to_data[image] = {
-                    "bboxes": [],
-                    "masks": [],
-                    "heat_map": [],
-                    "labels": [],
-                    "scores": [],
+        success, result = _EXECUTE.run_isolation(f"{code}\n{test}")
+        if verbosity == 2:
+            log_progress(
+                {
+                    "log": f"Debug attempt {count + 1}, reflection:",
+                    "result": fixed_code_and_test["reflections"],
                 }
-    return image_to_data
-
-
-def _handle_viz_tools(
-    image_to_data: Dict[str, Dict], tool_result: Dict
-) -> Dict[str, Dict]:
-    image_to_data = image_to_data.copy()
-
-    # handle grounding_sam_ and grounding_dino_
-    parameters = tool_result["parameters"]
-    # parameters can either be a dictionary or list, parameters can also be malformed
-    # becaus the LLM builds them
-    if isinstance(parameters, dict):
-        if "image" not in parameters:
-            return image_to_data
-        parameters = [parameters]
-    elif isinstance(tool_result["parameters"], list):
-        if len(tool_result["parameters"]) < 1 or (
-            "image" not in tool_result["parameters"][0]
-        ):
-            return image_to_data
-
-    for param, call_result in zip(parameters, tool_result["call_results"]):
-        # Calls can fail, so we need to check if the call was successful. It can either:
-        # 1. return a str or some error that's not a dictionary
-        # 2. return a dictionary but not have the necessary keys
-
-        if not isinstance(call_result, dict) or (
-            "bboxes" not in call_result
-            and "mask" not in call_result
-            and "heat_map" not in call_result
-        ):
-            return image_to_data
-
-        # if the call was successful, then we can add the image data
-        image = param["image"]
-        if image not in image_to_data:
-            image_to_data[image] = {
-                "bboxes": [],
-                "masks": [],
-                "heat_map": [],
-                "labels": [],
-                "scores": [],
-            }
+            )
+            _LOGGER.info(
+                f"Debug attempt {count + 1}, reflection: {fixed_code_and_test['reflections']}"
+            )
+            _CONSOLE.print(
+                Syntax(
+                    f"{code}\n{test}", "python", theme="gruvbox-dark", line_numbers=True
+                )
+            )
+            log_progress(
+                {
+                    "log": "Debug result:",
+                    "result": result,
+                }
+            )
+            _LOGGER.info(f"Debug result: {result}")
+        count += 1
 
-        image_to_data[image]["bboxes"].extend(call_result.get("bboxes", []))
-        image_to_data[image]["labels"].extend(call_result.get("labels", []))
-        image_to_data[image]["scores"].extend(call_result.get("scores", []))
-        image_to_data[image]["masks"].extend(call_result.get("masks", []))
-        # only single heatmap is returned
-        if "heat_map" in call_result:
-            image_to_data[image]["heat_map"].append(call_result["heat_map"])
-        if "mask_shape" in call_result:
-            image_to_data[image]["mask_shape"] = call_result["mask_shape"]
-
-    return image_to_data
-
-
-def sample_n_evenly_spaced(lst: Sequence, n: int) -> Sequence:
-    if n <= 0:
-        return []
-    elif len(lst) == 0:
-        return []
-    elif n == 1:
-        return [lst[0]]
-    elif n >= len(lst):
-        return lst
-
-    spacing = (len(lst) - 1) / (n - 1)
-    return [lst[round(spacing * i)] for i in range(n)]
-
-
-def visualize_result(all_tool_results: List[Dict]) -> Sequence[Union[str, Path]]:
-    image_to_data: Dict[str, Dict] = {}
-    for tool_result in all_tool_results:
-        # only handle bbox/mask tools or frame extraction
-        if tool_result["tool_name"] not in [
-            "grounding_sam_",
-            "grounding_dino_",
-            "extract_frames_",
-            "dinov_",
-            "zero_shot_counting_",
-            "visual_prompt_counting_",
-            "ocr_",
-        ]:
-            continue
-
-        if tool_result["tool_name"] == "extract_frames_":
-            image_to_data = _handle_extract_frames(image_to_data, tool_result)
-        else:
-            image_to_data = _handle_viz_tools(image_to_data, tool_result)
-
-    visualized_images = []
-    for image_str in image_to_data:
-        image_path = Path(image_str)
-        image_data = image_to_data[image_str]
-        if "_counting_" in tool_result["tool_name"]:
-            image = overlay_heat_map(image_path, image_data)
-        else:
-            image = overlay_masks(image_path, image_data)
-            image = overlay_bboxes(image, image_data)
-        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
-            image.save(f.name)
-            visualized_images.append(f.name)
-    return visualized_images
+    if verbosity >= 1:
+        _LOGGER.info("Final code and tests:")
+        _CONSOLE.print(
+            Syntax(f"{code}\n{test}", "python", theme="gruvbox-dark", line_numbers=True)
+        )
+        _LOGGER.info(f"Final Result: {result}")
+
+    return {
+        "code": code,
+        "test": test,
+        "success": success,
+        "test_result": result,
+        "working_memory": new_working_memory,
+    }
+
+
+def retrieve_tools(
+    plan: List[Dict[str, str]],
+    tool_recommender: Sim,
+    log_progress: Callable[[Dict[str, Any]], None],
+    verbosity: int = 0,
+) -> str:
+    tool_info = []
+    tool_desc = []
+    for task in plan:
+        tools = tool_recommender.top_k(task["instructions"], k=2, thresh=0.3)
+        tool_info.extend([e["doc"] for e in tools])
+        tool_desc.extend([e["desc"] for e in tools])
+    if verbosity == 2:
+        log_progress(
+            {
+                "log": "Retrieved tools:",
+                "tools": tool_desc,
+            }
+        )
+        _LOGGER.info(f"Tools: {tool_desc}")
+    tool_info_set = set(tool_info)
+    return "\n\n".join(tool_info_set)
 
 
 class VisionAgent(Agent):
-    r"""Vision Agent is an agent framework that utilizes tools as well as self
-    reflection to accomplish tasks, in particular vision tasks. Vision Agent is based
-    off of EasyTool https://arxiv.org/abs/2401.06201 and Reflexion
-    https://arxiv.org/abs/2303.11366 where it will attempt to complete a task and then
-    reflect on whether or not it was able to accomplish the task based off of the plan
-    and final results, if not it will redo the task with this newly added reflection.
+    """Vision Agent is an agentic framework that can output code based on a user
+    request. It can plan tasks, retrieve relevant tools, write code, write tests and
+    reflect on failed test cases to debug code. It is inspired by AgentCoder
+    https://arxiv.org/abs/2312.13010 and Data Interpeter
+    https://arxiv.org/abs/2402.18679
 
     Example
     -------
-        >>> from vision_agent.agent import VisionAgent
+        >>> from vision_agent import VisionAgent
         >>> agent = VisionAgent()
-        >>> resp = agent("If red tomatoes cost $5 each and yellow tomatoes cost $2.50 each, what is the total cost of all the tomatoes in the image?", image="tomatoes.jpg")
-        >>> print(resp)
-        "The total cost is $57.50."
+        >>> code = agent("What percentage of the area of the jar is filled with coffee beans?", media="jar.jpg")
     """
 
     def __init__(
         self,
-        task_model: Optional[Union[LLM, LMM]] = None,
-        answer_model: Optional[Union[LLM, LMM]] = None,
-        reflect_model: Optional[Union[LLM, LMM]] = None,
-        max_retries: int = 2,
-        verbose: bool = False,
+        planner: Optional[LLM] = None,
+        coder: Optional[LLM] = None,
+        tester: Optional[LLM] = None,
+        debugger: Optional[LLM] = None,
+        tool_recommender: Optional[Sim] = None,
+        verbosity: int = 0,
         report_progress_callback: Optional[Callable[[Dict[str, Any]], None]] = None,
-    ):
-        """VisionAgent constructor.
+    ) -> None:
+        """Initialize the Vision Agent.
 
         Parameters:
-            task_model: the model to use for task decomposition.
-            answer_model: the model to use for reasoning and concluding the answer.
-            reflect_model: the model to use for self reflection.
-            max_retries: maximum number of retries to attempt to complete the task.
-            verbose: whether to print more logs.
-            report_progress_callback: a callback to report the progress of the agent. This is useful for streaming logs in a web application where multiple VisionAgent instances are running in parallel. This callback ensures that the progress are not mixed up.
+            planner (Optional[LLM]): The planner model to use. Defaults to OpenAILLM.
+            coder (Optional[LLM]): The coder model to use. Defaults to OpenAILLM.
+            tester (Optional[LLM]): The tester model to use. Defaults to OpenAILLM.
+            debugger (Optional[LLM]): The debugger model to
+            tool_recommender (Optional[Sim]): The tool recommender model to use.
+            verbosity (int): The verbosity level of the agent. Defaults to 0. 2 is the
+                highest verbosity level which will output all intermediate debugging
+                code.
+            report_progress_callback: a callback to report the progress of the agent.
+                This is useful for streaming logs in a web application where multiple
+                VisionAgent instances are running in parallel. This callback ensures
+                that the progress are not mixed up.
         """
-        self.task_model = (
-            OpenAILLM(model_name="gpt-4-turbo", json_mode=True, temperature=0.0)
-            if task_model is None
-            else task_model
+
+        self.planner = (
+            OpenAILLM(temperature=0.0, json_mode=True) if planner is None else planner
         )
-        self.answer_model = (
-            OpenAILLM(model_name="gpt-4-turbo", temperature=0.0)
-            if answer_model is None
-            else answer_model
+        self.coder = OpenAILLM(temperature=0.0) if coder is None else coder
+        self.tester = OpenAILLM(temperature=0.0) if tester is None else tester
+        self.debugger = (
+            OpenAILLM(temperature=0.0, json_mode=True) if debugger is None else debugger
         )
-        self.reflect_model = (
-            OpenAILMM(model_name="gpt-4-turbo", json_mode=True, temperature=0.0)
-            if reflect_model is None
-            else reflect_model
+
+        self.tool_recommender = (
+            Sim(TOOLS_DF, sim_key="desc")
+            if tool_recommender is None
+            else tool_recommender
         )
-        self.max_retries = max_retries
-        self.tools = TOOLS
+        self.verbosity = verbosity
+        self.max_retries = 2
         self.report_progress_callback = report_progress_callback
-        if verbose:
-            _LOGGER.setLevel(logging.INFO)
 
     def __call__(
         self,
         input: Union[List[Dict[str, str]], str],
-        image: Optional[Union[str, Path]] = None,
-        reference_data: Optional[Dict[str, str]] = None,
-        visualize_output: Optional[bool] = False,
-        self_reflection: Optional[bool] = True,
+        media: Optional[Union[str, Path]] = None,
     ) -> str:
-        """Invoke the vision agent.
+        """Chat with Vision Agent and return intermediate information regarding the task.
 
         Parameters:
-            chat: A conversation in the format of
+            chat (List[Dict[str, str]]): A conversation in the format of
                 [{"role": "user", "content": "describe your task here..."}].
-            image: The input image referenced in the chat parameter.
-            reference_data: A dictionary containing the reference image, mask or bounding
-                box in the format of:
-                {"image": "image.jpg", "mask": "mask.jpg", "bbox": [0.1, 0.2, 0.1, 0.2]}
-                where the bounding box coordinates are normalized.
-            visualize_output: Whether to visualize the output.
-            self_reflection: boolean to enable and disable self reflection.
+            media (Optional[Union[str, Path]]): The media file to be used in the task.
+            self_reflection (bool): Whether to reflect on the task and debug the code.
 
         Returns:
-            The result of the vision agent in text.
+            str: The code output by the Vision Agent.
         """
+
         if isinstance(input, str):
             input = [{"role": "user", "content": input}]
-        return self.chat(
-            input,
-            image=image,
-            visualize_output=visualize_output,
-            reference_data=reference_data,
-            self_reflection=self_reflection,
-        )
-
-    def log_progress(self, data: Dict[str, Any]) -> None:
-        _LOGGER.info(data)
-        if self.report_progress_callback:
-            self.report_progress_callback(data)
-
-    def _report_visualization_via_callback(
-        self, images: Sequence[Union[str, Path]]
-    ) -> None:
-        """This is intended for streaming the visualization images via the callback to the client side."""
-        if self.report_progress_callback:
-            self.report_progress_callback({"log": "<VIZ>"})
-            if images:
-                for img in images:
-                    self.report_progress_callback(
-                        {"log": f"<IMG>base:64{convert_to_b64(img)}</IMG>"}
-                    )
-            self.report_progress_callback({"log": "</VIZ>"})
+        results = self.chat_with_workflow(input, media)
+        results.pop("working_memory")
+        return results  # type: ignore
 
     def chat_with_workflow(
         self,
         chat: List[Dict[str, str]],
-        image: Optional[Union[str, Path]] = None,
-        reference_data: Optional[Dict[str, str]] = None,
-        visualize_output: Optional[bool] = False,
-        self_reflection: Optional[bool] = True,
-    ) -> Tuple[str, List[Dict]]:
-        """Chat with the vision agent and return the final answer and all tool results.
+        media: Optional[Union[str, Path]] = None,
+        self_reflection: bool = False,
+    ) -> Dict[str, Any]:
+        """Chat with Vision Agent and return intermediate information regarding the task.
 
         Parameters:
-            chat: A conversation in the format of
+            chat (List[Dict[str, str]]): A conversation in the format of
                 [{"role": "user", "content": "describe your task here..."}].
-            image: The input image referenced in the chat parameter.
-            reference_data: A dictionary containing the reference image, mask or bounding
-                box in the format of:
-                {"image": "image.jpg", "mask": "mask.jpg", "bbox": [0.1, 0.2, 0.1, 0.2]}
-                where the bounding box coordinates are normalized.
-            visualize_output: Whether to visualize the output.
-            self_reflection: boolean to enable and disable self reflection.
+            media (Optional[Union[str, Path]]): The media file to be used in the task.
+            self_reflection (bool): Whether to reflect on the task and debug the code.
 
         Returns:
-            A tuple where the first item is the final answer and the second item is a
-            list of all the tool results. The last item in the tool results also
-            contains the visualized output.
+            Dict[str, Any]: A dictionary containing the code, test, test result, plan,
+                and working memory of the agent.
         """
-        if len(chat) == 0:
-            raise ValueError("Input cannot be empty.")
-
-        question = chat[0]["content"]
-        if image:
-            question += f" Image name: {image}"
-        if reference_data:
-            question += (
-                f" Reference image: {reference_data['image']}"
-                if "image" in reference_data
-                else ""
-            )
-            question += (
-                f" Reference mask: {reference_data['mask']}"
-                if "mask" in reference_data
-                else ""
-            )
-            question += (
-                f" Reference bbox: {reference_data['bbox']}"
-                if "bbox" in reference_data
-                else ""
-            )
-
-        reflections = ""
-        final_answer = ""
-        all_tool_results: List[Dict] = []
 
-        for _ in range(self.max_retries):
-            task_list = self.create_tasks(
-                self.task_model, question, self.tools, reflections
+        if len(chat) == 0:
+            raise ValueError("Chat cannot be empty.")
+
+        if media is not None:
+            for chat_i in chat:
+                if chat_i["role"] == "user":
+                    chat_i["content"] += f" Image name {media}"
+
+        code = ""
+        test = ""
+        working_memory: List[Dict[str, str]] = []
+        results = {"code": "", "test": "", "plan": []}
+        plan = []
+        success = False
+        retries = 0
+
+        while not success and retries < self.max_retries:
+            plan_i = write_plan(
+                chat,
+                TOOL_DESCRIPTIONS,
+                format_memory(working_memory),
+                self.planner,
+                media=[media] if media else None,
             )
-
-            task_depend = {"Original Question": question}
-            previous_log = ""
-            answers = []
-            for task in task_list:
-                task_depend[task["id"]] = {"task": task["task"], "answer": "", "call_result": ""}  # type: ignore
-            all_tool_results = []
-
-            for task in task_list:
-                task_str = task["task"]
-                previous_log = str(task_depend)
-                tool_results, call_results = self.retrieval(
-                    self.task_model,
-                    task_str,
-                    self.tools,
-                    previous_log,
-                    reflections,
-                )
-                answer = answer_generate(
-                    self.answer_model, task_str, call_results, previous_log, reflections
+            plan_i_str = "\n-".join([e["instructions"] for e in plan_i])
+            if self.verbosity >= 1:
+                self.log_progress(
+                    {
+                        "log": "Going to run the following plan(s) in sequence:\n",
+                        "plan": plan_i,
+                    }
                 )
 
-                tool_results["answer"] = answer
-                all_tool_results.append(tool_results)
+                _LOGGER.info(
+                    f"""
+{tabulate(tabular_data=plan_i, headers="keys", tablefmt="mixed_grid", maxcolwidths=_MAX_TABULATE_COL_WIDTH)}"""
+                )
 
-                self.log_progress({"log": f"\tCall Result: {call_results}"})
-                self.log_progress({"log": f"\tAnswer: {answer}"})
-                answers.append({"task": task_str, "answer": answer})
-                task_depend[task["id"]]["answer"] = answer  # type: ignore
-                task_depend[task["id"]]["call_result"] = call_results  # type: ignore
-            final_answer = answer_summarize(
-                self.answer_model, question, answers, reflections
+            tool_info = retrieve_tools(
+                plan_i,
+                self.tool_recommender,
+                self.log_progress,
+                self.verbosity,
             )
-            visualized_output = visualize_result(all_tool_results)
-            all_tool_results.append({"visualized_output": visualized_output})
-            if len(visualized_output) > 0:
-                reflection_images = sample_n_evenly_spaced(visualized_output, 3)
-            elif image is not None:
-                reflection_images = [image]
-            else:
-                reflection_images = None
+            results = write_and_test_code(
+                FULL_TASK.format(user_request=chat[0]["content"], subtasks=plan_i_str),
+                tool_info,
+                UTILITIES_DOCSTRING,
+                format_memory(working_memory),
+                self.coder,
+                self.tester,
+                self.debugger,
+                self.log_progress,
+                verbosity=self.verbosity,
+                input_media=media,
+            )
+            success = cast(bool, results["success"])
+            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})
 
             if self_reflection:
-                reflection = self_reflect(
-                    self.reflect_model,
-                    question,
-                    self.tools,
-                    all_tool_results,
-                    final_answer,
-                    reflection_images,
-                )
-                self.log_progress({"log": f"Reflection: {reflection}"})
-                parsed_reflection = parse_reflect(reflection)
-                if parsed_reflection["Finish"]:
-                    break
-                else:
-                    reflections += "\n" + parsed_reflection["Reflection"]
-            else:
-                self.log_progress(
-                    {"log": "Self Reflection skipped based on user request."}
+                reflection = reflect(
+                    chat,
+                    FULL_TASK.format(
+                        user_request=chat[0]["content"], subtasks=plan_i_str
+                    ),
+                    code,
+                    self.planner,
                 )
-                break
-        # '<ANSWER>' is a symbol to indicate the end of the chat, which is useful for streaming logs.
-        self.log_progress(
-            {
-                "log": f"The Vision Agent has concluded this chat. <ANSWER>{final_answer}</ANSWER>"
-            }
-        )
-
-        if visualize_output:
-            viz_images: Sequence[Union[str, Path]] = all_tool_results[-1][
-                "visualized_output"
-            ]
-            self._report_visualization_via_callback(viz_images)
-            for img in viz_images:
-                Image.open(img).show()
-
-        return final_answer, all_tool_results
-
-    def chat(
-        self,
-        chat: List[Dict[str, str]],
-        image: Optional[Union[str, Path]] = None,
-        reference_data: Optional[Dict[str, str]] = None,
-        visualize_output: Optional[bool] = False,
-        self_reflection: Optional[bool] = True,
-    ) -> str:
-        answer, _ = self.chat_with_workflow(
-            chat,
-            image=image,
-            visualize_output=visualize_output,
-            reference_data=reference_data,
-            self_reflection=self_reflection,
-        )
-        return answer
-
-    def retrieval(
-        self,
-        model: Union[LLM, LMM, Agent],
-        question: str,
-        tools: Dict[int, Any],
-        previous_log: str,
-        reflections: str,
-    ) -> Tuple[Dict, str]:
-        tool_id = choose_tool(
-            model,
-            question,
-            {k: v["description"] for k, v in tools.items()},
-            reflections,
-        )
-        if tool_id is None:
-            return {}, ""
-
-        tool_instructions = tools[tool_id]
-        tool_usage = tool_instructions["usage"]
-        tool_name = tool_instructions["name"]
+                if self.verbosity > 0:
+                    self.log_progress(
+                        {
+                            "log": "Reflection:",
+                            "reflection": reflection,
+                        }
+                    )
+                    _LOGGER.info(f"Reflection: {reflection}")
+                feedback = cast(str, reflection["feedback"])
+                success = cast(bool, reflection["success"])
+                working_memory.append({"code": f"{code}\n{test}", "feedback": feedback})
 
-        parameters = choose_parameter(
-            model, question, tool_usage, previous_log, reflections
-        )
-        if parameters is None:
-            return {}, ""
-        tool_results = {
-            "task": question,
-            "tool_name": tool_name,
-            "parameters": parameters,
-        }
+            retries += 1
 
         self.log_progress(
             {
-                "log": f"""Going to run the following tool(s) in sequence:
-{tabulate(tabular_data=[tool_results], headers="keys", tablefmt="mixed_grid", maxcolwidths=_MAX_TABULATE_COL_WIDTH)}"""
+                "log": f"Vision Agent has concluded this chat.\nSuccess: {success}",
+                "finished": True,
             }
         )
 
-        def parse_tool_results(result: Dict[str, Union[Dict, List]]) -> Any:
-            call_results: List[Any] = []
-            if isinstance(result["parameters"], Dict):
-                call_results.append(
-                    function_call(tools[tool_id]["class"], result["parameters"])
-                )
-            elif isinstance(result["parameters"], List):
-                for parameters in result["parameters"]:
-                    call_results.append(
-                        function_call(tools[tool_id]["class"], parameters)
-                    )
-            return call_results
-
-        call_results = parse_tool_results(tool_results)
-        tool_results["call_results"] = call_results
-
-        call_results_str = str(call_results)
-        return tool_results, call_results_str
+        return {
+            "code": code,
+            "test": test,
+            "test_result": results["test_result"],
+            "plan": plan,
+            "working_memory": working_memory,
+        }
 
-    def create_tasks(
-        self,
-        task_model: Union[LLM, LMM],
-        question: str,
-        tools: Dict[int, Any],
-        reflections: str,
-    ) -> List[Dict]:
-        tasks = task_decompose(
-            task_model,
-            question,
-            {k: v["description"] for k, v in tools.items()},
-            reflections,
-        )
-        if tasks is not None:
-            task_list = [{"task": task, "id": i + 1} for i, task in enumerate(tasks)]
-            task_list = task_topology(task_model, question, task_list)
-            try:
-                task_list = topological_sort(task_list)
-            except Exception:
-                _LOGGER.error(f"Failed topological_sort on: {task_list}")
-        else:
-            task_list = []
-        self.log_progress(
-            {
-                "log": "Planned tasks:",
-                "plan": task_list,
-            }
-        )
-        return task_list
+    def log_progress(self, data: Dict[str, Any]) -> None:
+        if self.report_progress_callback is not None:
+            self.report_progress_callback(data)
+        pass