"# AgentOptimizer: An Agentic Way to Train Your LLM Agent\n",
"\n",
"AutoGen offers conversable agents powered by LLM, tool, or human, which can be used to perform tasks collectively via automated chat. This framework allows tool use and human participation through multi-agent conversation.\n",
"Please find documentation about this feature [here](https://microsoft.github.io/autogen/docs/Use-Cases/agent_chat).\n",
"\n",
"In traditional ML pipeline, we train a model by updating its parameter according to the loss on the training set, while in the era of LLM agents, how should we train an agent? Here, we take an initial step towards the agent training. Inspired by the [function calling](https://platform.openai.com/docs/guides/function-calling) capabilities provided by OpenAI, we draw an analogy between model parameters and agent functions/skills, and update agent’s functions/skills based on its historical performance on the training set. As an agentic way of training an agent, our approach help enhance the agents’ abilities without requiring access to the LLMs parameters.\n",
"\n",
"In this notebook, we introduce a new class, ‘AgentOptimizer’, which is able to improve the function list of one Assistant-UserProxy pair according to the historical conversation histories. This feature would support agents in improving their ability to solve problems of the same type as previous tasks.\n",
"Specifically, given a set of training data, AgentOptimizer would iteratively prompt the LLM to optimize the existing function list of the AssistantAgent and UserProxyAgent with code implementation if necessary.\n",
"In the example scenario, we test the proposed AgentOptimizer in solving problems from the [MATH dataset](https://github.com/hendrycks/math). \n",
"AgentOptimizer is a class that is designed to improve the agents through optimizing its function call. It contains two core methods:\n",
"\n",
"1. `step()`: `step()` has three inputs: previous conversation history (history), the statistical information of solving previous problems (statistic), and the signature of current functions (func_signature). The output is a series of actions to manipulate the current functions.\n",
"\n",
"2. `update_function_call()`: This method updates the functions registered in the agents according to the actions from `step()`. "
" OPT_PROMPT = \"\"\"You are a function optimizer. Your task is to maintain a list of functions for the assistant according to the existing function list and conversation history that happens between the assistant and the user.\n",
"You can perform one of the following four actions to manipulate the function list using the functions you have:\n",
"1. Revise one existing function (using revise_function).\n",
"2. Remove one existing function (using remove_function).\n",
"3. Add one new function (using add_function).\n",
"4. Directly return \"TERMINATE\" to me if no more actions are needed for the current function list.\n",
"\n",
"Below are the principles that you need to follow for taking these four actions.\n",
"(1) Revise one existing function:\n",
"1. Pay more attention to the failed tasks and corresponding error information, and optimize the function used in these tasks according to the conversation history if needed.\n",
"2. A failed function call can occur due to incorrect input arguments (missing arguments) or an incorrect function code implementation. You should focus more on the function code implementation and make it easy to get success function call.\n",
"3. Do not revise the function that you think works well and plays a critical role in solving the problems according to the conversation history. Only making revisions if needed.\n",
"4. Sometimes, a NameError may occur. To fix this error, you can either revise the name of the function in the code implementation or revise the name of the function call to make these two names consistent.\n",
"1. The added new function should solve a higher-level question that encompasses the original query and extend the code's functionality to make it more versatile and widely applicable.\n",
"2. The added new function should solve queries of the same type, based on common reasoning steps without mentioning specific object names or entity terms.\n",
"3. Name the function and write the description concerning both the core reasoning pattern and data organization format, without referencing specific objects. The name of the function MUST be the same with the function name in the code you generated.\n",
"4. Replace specific strings or variable names with general variables to enhance the tool's applicability to various queries. All names used inside the function should be passed in as arguments.\n",
"If you think there is no need to perform any other actions for the current function list since the current list is optimal more actions will harm the performance in future tasks. Please directly reply to me with \"TERMINATE\".\n",
"\n",
"One function signature includes the following five elements:\n",
"The following table shows the statistical information for solving each task in each conversation and indicates whether each task was successfully solved.\n",
"According to the information I provide, please take one of four actions to manipulate list B using the functions you know.\n",
" \"\"\"\n",
"\n",
" ADD_FUNC = {\n",
" \"type\": \"function\",\n",
" \"function\": {\n",
" \"name\": \"add_function\",\n",
" \"description\": \"Add a function in the context of the conversation. Necessary Python packages must be declared. The name of the function MUST be the same with the function name in the code you generated.\",\n",
" \"description\": 'JSON schema of arguments encoded as a string. Please note that the JSON schema only supports specific types including string, integer, object, array, boolean. (do not have float type) For example: { \"url\": { \"type\": \"string\", \"description\": \"The URL\", }}. Please avoid the error \\'array schema missing items\\' when using array type.',\n",
" \"description\": \"A list of package names imported by the function, and that need to be installed with pip prior to invoking the function. This solves ModuleNotFoundError. It should be string, not list.\",\n",
" \"description\": \"Revise a function in the context of the conversation. Necessary Python packages must be declared. The name of the function MUST be the same with the function name in the code you generated.\",\n",
" \"description\": 'JSON schema of arguments encoded as a string. Please note that the JSON schema only supports specific types including string, integer, object, array, boolean. (do not have float type) For example: { \"url\": { \"type\": \"string\", \"description\": \"The URL\", }}. Please avoid the error \\'array schema missing items\\' when using array type.',\n",
" \"description\": \"A list of package names imported by the function, and that need to be installed with pip prior to invoking the function. This solves ModuleNotFoundError. It should be string, not list.\",\n",
" \"description\": \"Remove one function in the context of the conversation. Once remove one function, the assistant will not use this function in future conversation.\",\n",
"This agent is a customozied MathUserProxy inherits from its [partent class](https://github.com/microsoft/autogen/blob/main/autogen/agentchat/contrib/math_user_proxy_agent.py.) \n",
"\n",
"It supports using both function_call and python to solve math problems."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def is_termination_msg_mathchat(message):\n",
" if isinstance(message, dict):\n",
" message = message.get(\"content\")\n",
" if message is None:\n",
" return False\n",
" cb = extract_code(message)\n",
" contain_code = False\n",
" for c in cb:\n",
" if c[0] == \"python\" or c[0] == \"wolfram\":\n",
" contain_code = True\n",
" break\n",
" if message.rstrip().find(\"TERMINATE\") >= 0:\n",
" DEFAULT_REPLY = \"Continue. Please keep solving the problem until you need to query. (If you get to the answer, put it in \\\\boxed{}.)\"\n",
" PROMPTS = \"\"\"Let's solve a math problem.\n",
"Query requirements:\n",
"You should always use the 'print' function for the output and use fractions/radical forms instead of decimals.\n",
"You can use packages like sympy to help you.\n",
"You must follow the formats below to write your code:\n",
"```python\n",
"# your code\n",
"```\n",
"If some packages are missing, you could also suggest a code to install the corresponding package.\n",
"\n",
"Please follow this process:\n",
"1. Solve the problem step by step (do not over-divide the steps).\n",
"2. Take out any queries that can be asked through Python code (for example, any calculations or equations that can be calculated) and functions you know in the context of this conversation.\n",
" revise_prompt = \"You may make a wrong function call (It may due the arguments you provided doesn't fit the function arguments like missing required positional argument). \\\n",
" If you think this error occurs due to you make a wrong function arguments input and you could make it success, please try to call this function again using the correct arguments. \\\n",
" Otherwise, the error may be caused by the function itself. Please directly reply me with TERMINATE. We need to terminate the conversation. \"\n",
" return True, \"The result is Correct. Please reply me with TERMINATE.\"\n",
" else:\n",
" self.logs[\"is_correct\"] = 0\n",
" return False, None\n",
" else:\n",
" return False, None\n",
"\n",
" def _reset(self):\n",
" self._valid_q_count = 0\n",
" self._total_q_count = 0\n",
" self._accum_invalid_q_per_step = 0\n",
" self._previous_code = \"\"\n",
" self.last_reply = None\n",
"\n",
" self.query = None\n",
" self.answer = None\n",
" self.logs = {}"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Load dataset\n",
"\n",
"MATAH dataset contains 12,500 challenging competition mathematics problems. Each problem in MATH has a full step-by-step solution which can be used to teach models to generate answer derivations and explanations. \n",
"\n",
"We strctly follow the train/test splits of [Craft](https://github.com/lifan-yuan/CRAFT). Please specific your own path to the dataset. Here we sample the first 10 algebra problems as examples. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"test_data, train_data = [], []\n",
"with open(\"MATH/dataset/algebra.jsonl\", \"r\", encoding=\"utf-8\") as f:\n",
" for line in f:\n",
" test_data.append(json.loads(line))\n",
"with open(\"MATH/dataset/train/algebra.jsonl\", \"r\", encoding=\"utf-8\") as f:\n",
"success_rate_without_agent_training = sum / len(statistic_list.keys())"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Agent optimizing \n",
"\n",
"Then, use the AgentOptimizer to iteratively optimize the agents by optimizing the function calls according to the historical conversations and performance. \n",
"\n",
"Here we optimize these two agents for ten epochs. "
"{'name': 'evaluate_expression', 'description': 'Evaluate arithmetic or mathematical expressions provided as strings.', 'arguments': {'expression': {'type': 'string', 'description': 'The mathematical expression to evaluate.'}}, 'packages': 'sympy', 'code': 'from sympy import sympify, SympifyError\\n\\ndef evaluate_expression(expression):\\n try:\\n result = sympify(expression)\\n if result.is_number:\\n result = float(result)\\n else:\\n result = str(result)\\n return result\\n except SympifyError as e:\\n return str(e)'}\n",
"{'name': 'calculate_compound_interest_principal', 'description': 'Calculate the principal amount needed to achieve a certain future value with quarterly compound interest.', 'arguments': {'future_value': {'type': 'number', 'description': 'The total amount of money desired in the future.'}, 'annual_interest_rate': {'type': 'number', 'description': 'The annual interest rate in decimal form.'}, 'compounding_periods': {'type': 'integer', 'description': 'The number of times interest is compounded per year.'}, 'years': {'type': 'number', 'description': 'The time in years the money is invested for.'}}, 'packages': 'sympy', 'code': \"from sympy import symbols, solve, N\\n\\nfuture_value, annual_interest_rate, compounding_periods, years = symbols('future_value annual_interest_rate compounding_periods years')\\nP = symbols('P')\\nequation = Eq(future_value, P * (1 + annual_interest_rate / compounding_periods) ** (compounding_periods * years))\\nprincipal = solve(equation, P)[0]\\nprincipal = N(principal, chop=True).evalf()\"}\n",
"{'name': 'solve_linear_system', 'description': 'Solve a system of linear equations represented as coefficients and variables.', 'arguments': {'equations': {'type': 'array', 'items': {'type': 'array', 'items': {'type': 'number'}}}, 'variables': {'type': 'array', 'items': {'type': 'string'}}}, 'packages': 'sympy', 'code': 'from sympy import Matrix, symbols, linsolve\\n\\ndef solve_linear_system(equations, variables):\\n if not equations or not all(len(eq) == len(variables) + 1 for eq in equations):\\n return \"Error: Equations list is empty or not all equations have the correct length.\"\\n if len(variables) != len(set(variables)):\\n return \"Error: Duplicate symbols in the \\'variables\\' list.\"\\n try:\\n sym_vars = symbols(\\' \\'.join(variables))\\n matrix = Matrix(equations)\\n system = (matrix, sym_vars)\\n solution = linsolve(system)\\n solution_list = list(solution)\\n if not solution_list:\\n return \"Error: No solution exists for the given system of equations.\"\\n result = solution_list[0] if solution_list else []\\n return dict(zip(variables, result))\\n except Exception as e:\\n return f\"Error: {str(e)}\"'}\n",
