TaskWeaver: A Code-First Agent Framework

We introduce some important concepts in the TaskWeaver system.

Python has emerged as the de-facto language for data analysis, and TaskWeaver leverages this popularity by converting user requests into Python programs that run on dedicated processes. Users can issue natural language requests to manipulate the state of the Python process with TaskWeaver, which utilizes popular libraries like numpy, pandas, sklearn, etc. Unlike other frameworks that rely on text or file-based expressions, TaskWeaver utilizes native data structures such as pandas DataFrame. This makes it easy to perform tasks such as pulling data from a database, running machine learning algorithms (e.g., anomaly detection, classification or clustering), summarizing results, and visualizing analysis outcomes.

The TaskWeaver approach involves a process of code verification (before code execution) to ensure that the generated code conforms to pre-defined rules and does not violate any safety requirements. This is achieved through the use of configurable rules that can be customized to suit different scenarios. By doing so, users can ensure that the code does not import any forbidden packages or invoke any unsafe functions. TaskWeaver also supports a dedicated plugin-only mode, leveraging the function calling capability of LLMs, where only the plugin functions can be called, and all other generated code is denied.

When users make ad-hoc requests for data analysis, it often involves multiple iterations. As a result, TaskWeaver needs to maintain the state of code execution throughout the entire session. This is similar to programming in Python using a Jupyter Notebook, where users type code snippets in a sequence of cells and the program’s internal state progresses sequentially. However, in TaskWeaver, users use natural language instead of a programming language. TaskWeaver converts each user request into one or more code snippets in each round, depending on the specific plan.

The Planner decomposes users’ requests into sub-tasks that are simplified and executable by the CI. We identified three types of dependencies between different sub-tasks: sequential, interactive, and none. Sub-task B is sequentially dependent on sub-task A means that B must run after A is done. Interactive dependency is the same as sequential in the execution order. More than that, it means there is a need for intervention by a human or a LLM between the two sub-tasks. For example, the user instructs TaskWeaver to read a file and follow the instructions in its content. The Planner split the request into 2 steps: (1) read the file content; (2) follow the instructions. The LLM needs to complete the first step to read the content before carrying out the second step. The third type of dependency is none which means there is no dependency there, and they can be conducted in parallel. According to this definition, two steps where one sequentially depends on the other could be merged into one step for code generation because no intervention is required. In TaskWeaver, the Planner involves two-phase planning, where we enforce it to reconsider the generated plan of sub-tasks to merge sub-tasks with sequential dependency. This can prevent the model from splitting the request into too fine-grained steps leading to a prolonged execution process and incurring unnecessary costs by calling the LLM many times.

TaskWeaver is designed with the capacity to rectify errors throughout the planning and code generation stages. When it detects that the outcomes of the preceding steps diverge from the anticipated results, the Planner possesses the capability to reassess and modify its plan, exploring alternative approaches. Furthermore, the CI is equipped to evaluate the results of code execution. Should the code execution process encounter an exception or the code fails to pass the verification, the CI can initiate a re-generation of the code, thereby attempting to correct the code. This iterative process ensures that TaskWeaver maintains a robust and adaptable approach to task execution, increasing the reliability and efficiency of the framework.

In TaskWeaver, plugins are specialized Python functions used to handle tasks that are either too complex or require specific domain knowledge, thereby eliminating the number of plugins needed since TaskWeaver can already handle general Python code generation. Furthermore, TaskWeaver features dynamic plugin selection. This means that after a user request is received, only the plugins that are relevant to that request are selected from a pool of available plugins. This approach ensures that TaskWeaver uses the most appropriate tools for the task at hand without overloading the prompt with unnecessary functions.

One way to incorporate domain knowledge is by defining custom plugins, as discussed previously. However, for domain-specific tasks, it can be challenging for the LLM to generate the correct code to call the plugins or to make a good plan. To address this, we have introduced an interface in TaskWeaver that allows users to guide the system for such difficult tasks. TaskWeaver enables users to configure examples to teach the LLM how to respond to certain requests. For instance, a conversation history containing step-by-step thoughts and request/response sections can be used as an example. There are two types of examples in TaskWeaver: one is used for planning and the other for code generation. By using examples, TaskWeaver can incorporate domain-specific knowledge and improve the LLM’s ability to generate accurate plans and code for difficult tasks.

Executing freely generated code can introduce security risks. For example, a malicious user may ask TaskWeaver to open a file containing security keys, delete system files, or terminate a process. A typical way to mitigate the security risks is to include certain statements in the prompts to the LLM, which is implemented in TaskWeaver. To further prevent these malicious behaviors, TaskWeaver allocates a separate worker process running inside a session-associated docker container. This architectural choice ensures that each worker process operates independently, isolated from processes of the host OS as well as other sessions, thus effectively preventing malicious behaviors.

Existing LLM frameworks such as Langchain [5] make it easy to build proof-of-concept demos. However, building a reliable system with these frameworks can be time-consuming due to the large number of components they provide (like a large box of Lego bricks). One of the main goals of TaskWeaver is to make it easy to use. Typically, users only need to customize the plugins to get started. For more difficult tasks, users can customize the examples for code generation and planning. To aid users, TaskWeaver includes various tools to help, such as a tool to convert Python functions into plugins and a tool to save existing conversations as examples. These tools make it easier for users to customize the system. Another key feature of TaskWeaver is its support for multi-tenant. The system implements a session manager to isolate different user sessions, making it easy to serve multiple users as a service.

The cost of calling LLMs can be significant, and TaskWeaver addresses this issue by letting different roles be configured with different LLM models. For instance, GPT 3.5 is much cheaper compared to GPT 4. Therefore, for simpler tasks, we may use a cheaper model to reduce the overall cost.

The Planner serves as a crucial component in the TaskWeaver system, where it decomposes requests from users into several sub-tasks and orchestrates capabilities within TaskWeaver to complete the task and report back to the users. From a high-level perspective, the Planner functions as the entry point and controller, managing the entire system. As illustrated in Figure 2, the Planner communicates bidirectionally with each component, sending queries and receiving responses.

As shown in Fig. 3, a typical working routine of the Planner is as follows: The Planner first receives a query from users and then decides to decompose it into multiple sub-tasks. These sub-tasks essentially form a “Initial Plan”, which is generated based on the knowledge of LLMs or enhanced by domain-specific “Examples”. After drafting the initial plan, the Planner is requested to refine this initial plan by considering the dependencies among the sub-tasks as discussed in Sec. 4.4 in a chain-of-thought manner. The Planner may merge multiple sub-tasks into one in its “Final Plan”. Table 1 shows two examples of the Planner’s initial plan and final plan w.r.t the User’s request. In the first example, the initial plan has 4 steps while the first 3 steps are merged into one in the final plan. In the second example, the final plan is the same as the initial plan because the Planner needs to read the file content and then understand what would be the next step.

After the plan is finalized, the Planner takes action by assigning each sub-task, with phrased queries, to ask the CI to generate a code snippet, leveraging in-domain plugins when necessary. The execution results are sent back to the Planner. Following the ReAct (reasoning and act) design pattern, upon observing the execution results, the Planner may update its plan, request additional information from users, and so on. The process is repeated for subsequent sub-tasks until the entire plan is completed.

When addressing domain-specific scenarios where the LLM’s own knowledge is insufficient, the Planner can incorporate external knowledge with “Examples”. These examples are tailored by scenario developers according to their usage requirements. One example is essentially the chat history between the User, the Planner, and the Code Interpreter, including the plans of the Planner. The Examples could be saved from an online conversation for the purpose of stabilizing the planning process for future requests, or they could be manually prepared. A real planning Example is provided in Appendix B.

Code Generator is designed to utilize LLMs to automatically synthesize a Python code snippet based on an incoming request. From a high-level view, CG combines the benefits of both the plugin system and code interpreter, allowing the capability of invoking plugins and generating additional code. Inside the CG, plugins and examples customized by users are leveraged to generate the code.

1. 1.

   Plugin: In general, plugins can take on various forms such as a web API call, a software module, a customized algorithm, or a deep learning model. Regardless of the form, all variants can be invoked by a function call. Therefore, we have encapsulated the plugins as a Python function within the Code Generator for seamless invocation.

2. 2.

   Example: To aid the LLMs in adapting to domain-specific scenarios, examples are designed to help guide LLMs to behave by following the examples. Examples act as a guide for the LLMs to follow and improve their performance by providing contextualized examples. This feature ensures that the generated code aligns with the specific requirements of the user’s domain. A real code generation Example is provided in Appendix C.

The CG has the ability to generate code that exclusively calls upon plugins, code that does not use any plugins, or a combination of both. The code that doesn’t depend on plugins is specifically tailored to handle requests that cannot be met using only existing plugins. In the motivating example presented in Section 2, there is a plugin called ‘anomaly_detection’ that can identify anomalies in data. In this scenario, the CG would first invoke the plugin through a single line of code. Subsequently, the CG would generate multiple lines of additional code to visualize the detected anomalies, as requested by the user.

The knowledge of LLMs is often limited, particularly in domain-specific scenarios. In many cases, there are in-domain tools better suited for completing a specific task. Therefore, the Plugin system can enhance LLMs by enabling them to leverage existing tools. The Plugin system consists of the following two components:

1. 1.

   Plugin Schema LLMs need to understand the capabilities of a plugin, including its arguments and return values. This information is embedded in the plugin schema file (in YAML format) and could be customized by users. The plugin schema comprises the name, metadata, plugin description, arguments, and return values. For each argument and return value, users should provide its name, type, and description. The plugin schema is then supplied to LLMs as part of the prompt to generate the code.

2. 2.

   Plugin Implementation The plugin implementation defines how each plugin is executed using Python in Code Executor, where the arguments and return values should correspond with the Plugin Schema. Various plugins can be implemented in Python code, such as a domain-specific algorithm, a software module, a deep learning model or a Web API call. It is important to note that LLMs do not need to know the plugin implementation details to generate code. On the contrary, the Plugin Schema is used for code generation.

Appendix D shows an example of the plugin schema and implementation.

Even with explicit instructions added to the prompt, such as restricting the call to plugin functions, prohibiting local file system modifications, or disallowing certain package installations, the LLM may still generate code that disregards these instructions. This necessitates a post-verification process to ensure the generated code is safe to execute. Following code generation, the CG parses the code into an Abstract Syntax Tree (AST) and examines it line by line for any violations. If violations are detected, an error message is reported, and a new code snippet is generated. The CG will retry this process several times; if violations persist, the CG will notify the planner of its failure to generate compliant code.

Code auto-correction shares similarities with the post-verification process. If the generated code fails to execute, the exception message is reported to the CG, prompting it to retry generating the code. The primary difference is that post-verification operates within the CG, while code auto-correction relies on the CE to run the code and report errors. Notably, if the execution fails, the Planner can directly redirect the message to the CG instead of synthesizing the report via the LLM. TaskWeaver permits the code to regenerate up to a maximum of three times.

The Code Executor (CE) receives the code generated by the CG, collects dependent modules and plugin definitions, executes it, preserves context information such as logs, and returns an execution result to the Planner. In TaskWeaver, we implement CE based on the Python Jupyter kernel. To prevent interference between different sessions, CE maintains a separate Jupyter process for each session. In our current implementation, we support two execution modes: local and container. In the local mode, the Jupyter kernel is launched as a local process, while in the container model, the kernel is running inside a Docker container. When a session concludes, the Jupyter process is terminated. The execution result contains the following parts which are returned to the planner.

• •

  Return Code: a successful code execution will return code 0 and otherwise 1.

• •

  Logs: Logs are generated in two ways: (1) the stdout/stderr output of the program, and (2) log messages recorded using a logging utility within the plugins.

• •

  Output: This is the output of the Jupyter cell running the generated code.

• •

  Artifacts: The generated code or the plugin calls may produce artifacts such as a CSV file or an image. The user can download the artifacts via the provided URLs.

If the CE fails to execute the code from the CG, it will report the error logs to the CG, who will attempt to revise the code to fix the issue. All the error information and the failed code are kept in the CG’s conversation history so that the CG is aware of the full execution history and can precisely understand the state of the CE.

In practical scenarios, when a user tasks TaskWeaver with solving a complex problem, the system may initially falter. However, with repeated attempts and additional instructions from the user, TaskWeaver can eventually find a solution. A challenge arises when a user presents a similar or identical problem at a later time: TaskWeaver struggles to provide the correct solution promptly because it lacks the capability to remember past experiences. To address this, we have introduced a feature known as experience memory within the Memory module of TaskWeaver.

Through this mechanism, a user can command TaskWeaver to record the chat history. The system then distills ‘experience tips’ from this history, which encapsulate actionable insights about what to do—or what not to do—in response to requests akin to those encapsulated in the tips. These insights are then stored in an ‘experience pool’. When TaskWeaver encounters similar requests in the future, it will draw upon this pool to inform its strategy for planning and code generation. The retrieved experiences are integrated into the prompts used by the Planning and Code Generation (CG) components, thereby enhancing TaskWeaver’s response to subsequent, similar requests.

The utility of the experience memory extends beyond merely aiding in the resolution of challenging problems; it also serves to capture and incorporate user preferences into interactions with the agent. For instance, if a user exhibits a predilection for plotting charts with specific color schemes, this preference is gleaned from the chat history and preserved as an experience tip. Consequently, such personalized preferences are automatically considered in relevant future tasks.

We conducted evaluations on all the test cases from the datasets mentioned earlier. Each test case awards one or several points upon successful completion. As an agent provides solutions, we count the points earned and compute a normalized score ranging from 0 to 1 to represent the test case’s final score. We put TaskWeaver to the test using various LLMs as the underlying model. For each LLM, we performed at least two runs of all test cases and calculated the mean score, which we present as the LLM’s overall performance score. The summarized results are shown in the table below.

The performance of GPT-4 consistently surpasses that of GPT-3.5, underscoring its enhanced capabilities in comprehending problems, strategizing task execution, and generating code. When evaluated across three external benchmark datasets, both GPT-3.5 and GPT-4 achieve their highest scores on the InfiAgent-DABench. This superior performance can be attributed to TaskWeaver’s proficiency in tasks that involve data loading followed by analytical operations. However, when it comes to the DS-1000 dataset, which is derived from StackOverflow queries, the models face challenges. The ambiguity of the questions and the complexity of evaluating answers, particularly those involving DataFrames, make accurate assessment difficult. This suggests that while GPT-4’s advancements are notable, there is still room for improvement in processing and generating precise answers from less structured and clear data sources.

The detailed evaluation results for the 23 test cases are depicted in Table 4. Although these cases were crafted to assess compliance with our predefined requirements, it is important to note that their successful completion by TaskWeaver is not unequivocally assured. This is attributed to various factors, including the inherent variability in LLM performance. A notable observation from our study is that the performance of a given LLM endpoint is not consistent, exhibiting fluctuations over time and across different host machines. Such variability poses a significant challenge for applications based on LLMs.

In Table 4, the ‘Score’ column enumerates the points accrued across various test cases when utilizing GPT3.5 and GPT4. Meanwhile, the ‘Normalized Score’ column provides a proportional representation of these scores, scaled relative to the maximum possible points for each case. On average, the normalized scores for GPT3.5 and GPT4 stand at 0.42 and 0.85, respectively, indicating a substantial performance differential between the two models.

Upon detailed examination of Table 4, it becomes evident that GPT3.5 surpasses GPT4 in only two instances, specifically in test cases #9 and #16. In test case #9, which involves predicting stock prices, both GPT3.5 and GPT4 successfully downloaded historical price data. However, only GPT3.5 managed to accurately predict future prices within the evaluation period. For test case #16, which requires listing files in the current working directory, GPT4 did not execute the task, citing an inability to access the local file system, while GPT3.5 completed the task without issue.

Following the transformation process in Section 7.1, the DS-1000 dataset comprises 816 individual test cases. These cases are categorized across five widely-utilized libraries dedicated to data analytics tasks, with the distribution illustrated in Fig. 7. To modify the complexity of the original test cases, the developers of the DS-1000 have introduced various perturbations. The variety of these perturbation types and their respective frequencies within the test cases are depicted in Fig. 7.

Figure 9 presents the normalized scores for various libraries when processed by two different Large Language Models (LLMs), specifically GPT4 and GPT3.5. On average, GPT4 attains a normalized score of 0.6, surpassing the average score of 0.4 achieved by GPT3.5. An analysis of performance across the five libraries indicates a consistent pattern for both models, with Matplotlib recording the highest normalized score and Pandas the lowest. In light of the lower scores for Pandas, we conducted a closer examination of the cases where the models underperformed. A common issue identified was the challenge in accurately determining the correctness of Pandas test cases, which often require a comparison between two DataFrames. Such comparisons are prone to various errors, including discrepancies in indexes or data types. Nevertheless, these mismatches might not always signify actual failures in a practical context, as they are generally simple to rectify and might not pose a real problem.

Figure 9 displays the normalized scores associated with various perturbation types. It is observable that the performance of GPT3.5 deteriorates as the complexity of the perturbations escalates, which corroborates our initial predictions. Conversely, GPT4’s performance does not mirror that of GPT3.5, particularly in the context of Semantic perturbations. This divergence can be attributed to GPT4’s enhanced semantic comprehension abilities, which enable it to handle such perturbations more adeptly than GPT3.5, suggesting that these modifications do not pose the same level of difficulty for more advanced models.

All test cases in the InfiAgent-BABench are about analyzing the data loaded from a CSV file and then answering one or more questions. Figure 11 shows the distribution of the number of scoring points of the questions in this dataset. Most of the test cases have less than 3 questions. We conducted the evaluations with the test cases and collected the scoring points gained by the agent, one scoring point for each question. Then, we calculated the normalized score for each test case, i.e., a value between 0 to 1. Figure 11 shows the normalized scores of test cases with different scoring points. The average normalized scores of GPT4 and GPT3.5 are 0.88 and 0.70, respectively.

In this experiment, we conducted a comparative study of GPT-3.5 and GPT-4 across multiple datasets on DSEval benchmark, revealing significant enhancements in GPT-4’s performance. The %PASS denotes the pass rate where TaskWeaver’s execution result is exactly the same as the groundtruth one. Comparing the executions of two programs is often very hard. Sometimes, although the agent has fulfilled the user’s question, it does not imply that the final results are the same and comparable. Two common errors might lead to this situation, Presentation Errors (PE) and Intact Violation (IV). Presentation Errors manifest when the output is nearly correct but is flawed in terms of formatting or the method of presentation. For instance, an agent might neglect to capitalize a column name as required or mistakenly direct results to the console rather than placing them within the designated cell outputs. Intact Violations occur when the solution closely approaches correctness but fails to maintain the data’s original state, often seen when intermediate computations are necessary, and the agent alters the original dataset unnecessarily.

The comparative outcomes, as delineated in Table 5, underscore the substantial strides made by GPT-4 in diminishing the incidence of Presentation Errors (PE) and Intact Violations (IV) across all examined datasets. Notably, within the Leetcode dataset, there was a pronounced increase in the correctness score from 42.50 for GPT-3.5 to 75.00 for GPT-4. Moreover, there was a notable enhancement in the metric representing results unmarred by PE or IV, escalating from 60.00 to 77.50. This pattern of improvement was consistent across additional datasets, including Stack Overflow (SO). While the DSEval benchmark offers a framework for performance evaluation by furnishing ground truth code, it is not without limitations; the rigidity of automated comparisons can lead to undervaluation of the AI’s performance, as it may fail to recognize solutions that, from a human perspective, have effectively fulfilled the task requirements. This highlights the need for a more nuanced evaluation approach that can accurately reflect the practical utility of AI-generated code.

We aim to identify anomalies within a time series dataset stored in an SQL database (sqlite3). To accomplish this goal, we need to integrate the two plugins for data retrieval and time series anomaly detection:

• •

  pull_data_sql(nl_query) $\rightarrow$ sql, dataframe: Given a natural language query, pull data from a database and return the result in a pandas DataFrame, along with the generated SQL query.

• •

  anomaly_detection $\rightarrow$ dataframe with anomalies, results description: This plugin utilizes a straightforward 3-sigma algorithm to identify any abnormal data points in a time series. These data points deviate from the mean value by more than three standard deviations.

Note: There is no need to provide plugins for general tasks, such as reading or writing files, as the TaskWeaver is capable of generating the necessary code for these tasks. This is an advantage of TaskWeaver that reduces the development effort to bootstrap applications.

Fig.12 in Appendix A illustrated the detailed plan generated by the Planner of TaskWeaver:

1. 1.

   Instruct CodeInterpreter to pull data from the time_series table in the database.

2. 2.

   Confirm the columns to be detected anomalies.

3. 3.

   Instruct CodeInterpreter to detect anomalies on the pulled data.

4. 4.

   Report the detected anomalies to the user.

In the first step, the Code Interpreter invoked the pull_data_sql plugin function and automatically filled the required parameters.

After obtaining the time series data from the database, TaskWeaver asked the user to provide the column names as additional information to proceed with anomaly detection. The user then provided the "ts" and "val" column names to detect anomalies. With the additional information confirmed, TaskWeaver proceeded to execute the third step, i.e., detecting anomalies using the anomaly_detection plugin.

In the end, 11 anomalies were detected and a response was composed for the user.

The user requested TaskWeaver to forecast the Nasdaq 100 (QQQ) index prices for the next seven days with the ARIMA model. The detailed execution process is shown in Fig.13 in Appendix A. Here is the plan generated by the TaskWeaver Planner:

1. 1.

   Instruct CodeInterpreter to obtain historical QQQ stock price data.

2. 2.

   Preprocess the data.

3. 3.

   Train the ARIMA model.

4. 4.

   Forecast the stock price for the next 7 days.

5. 5.

   Report the predicted stock prices to the user.

After receiving the first subtask from the Planner, the Code Interpreter proceeded to generate a code snippet for the first step. The Code Interpreter thought to itself, "ProgramApe (internal code name) will generate a code snippet to load QQQ historical data using pandas_datareader library." The resulting code snippet is as follows:

Nonetheless, the execution encountered a failure and the following error information was reported:

As a result, the Code Interpreter had to modify the code to resolve this issue.

In the aforementioned code, we can observe that TaskWeaver employed an alternative approach to search for stock price data using the yfinance library. This time, the historical QQQ stock price data was successfully downloaded and the results were forwarded to the Planner for further instructions.

The Planner received the downloaded stock price data and proceeded with the following steps of preprocessing and ARIMA model training. Ultimately, TaskWeaver utilized the fitted ARIMA model to predict the stock prices for the next 7 days.