From 217321127f2557728deed2ebefe733fcca9fcb8c Mon Sep 17 00:00:00 2001 From: Misraaks <9754171+ampersand-five@users.noreply.github.com> Date: Fri, 5 Apr 2024 12:37:11 -0500 Subject: [PATCH] Issue #274: Fix LLMCompiler Bugs (#275) * Update LLMCompiler.ipynb The code didn't actually handle when there's more than one dependency listed in an argument. The example of multi-step math used in the code doesn't work because the third step has more than one dependency in its args and the code only handles one: ``` User query: "What's ((3*(4+5)/0.5)+3245) + 8? What's 32/4.23? What's the sum of those two values?" 1. math(problem="((3*(4+5)/0.5)+3245) + 8") 2. math(problem="32/4.23") 3. math(problem="${1}+${2}") 4. join() ``` Step 3 fails, and does so silently too. It correctly parses strings that are "${1}", but not multi-argument strings like "${1}+${2}" This fixes it to handle any number of dependencies listed in one arg. It uses the regex pattern already defined in output_parser.py. * Fix looping function calls because no arg context Task Fetching Unit, when it inserts the tool_messages, the function args don't get logged with function calls, this causes the replanner to loop and use the same function parameters and never corrects itself. * Fix tool index count and range In the planner prompt formatting, the 'num_tools' variable needs to have a +1 added to it because we're appending the join() function, without +1 it says there's one less available tools than there actually is because it only counts the passed in functions and not the join() function. And add +1 for listing otherwise it starts at a 0 offset. * Fix unhandled empty iterator exception Task fetching unit: When it calls the tools it doesn't handle when there's no tasks. next() is called and fails. This fixes it by wrapping in a try/except. On exception we set tasks to an empty list since next() failed, there's no tasks. --- examples/llm-compiler/LLMCompiler.ipynb | 42 ++++++++++++++++--------- 1 file changed, 28 insertions(+), 14 deletions(-) diff --git a/examples/llm-compiler/LLMCompiler.ipynb b/examples/llm-compiler/LLMCompiler.ipynb index e77331bda..24d9fcc52 100644 --- a/examples/llm-compiler/LLMCompiler.ipynb +++ b/examples/llm-compiler/LLMCompiler.ipynb @@ -222,11 +222,11 @@ " llm: BaseChatModel, tools: Sequence[BaseTool], base_prompt: ChatPromptTemplate\n", "):\n", " tool_descriptions = \"\\n\".join(\n", - " f\"{i}. {tool.description}\\n\" for i, tool in enumerate(tools)\n", + " f\"{i+1}. {tool.description}\\n\" for i, tool in enumerate(tools) # +1 to offset the 0 starting index, we want it count normally from 1.\n", " )\n", " planner_prompt = base_prompt.partial(\n", " replan=\"\",\n", - " num_tools=len(tools),\n", + " num_tools=len(tools)+1, # Add one because we're adding the join() tool at the end.\n", " tool_descriptions=tool_descriptions,\n", " )\n", " replanner_prompt = base_prompt.partial(\n", @@ -236,7 +236,7 @@ " ' - When starting the Current Plan, you should start with \"Thought\" that outlines the strategy for the next plan.\\n'\n", " \" - In the Current Plan, you should NEVER repeat the actions that are already executed in the Previous Plan.\\n\"\n", " \" - You must continue the task index from the end of the previous one. Do not repeat task indices.\",\n", - " num_tools=len(tools),\n", + " num_tools=len(tools)+1,\n", " tool_descriptions=tool_descriptions,\n", " )\n", "\n", @@ -338,6 +338,7 @@ "source": [ "from typing import Any, Union, Iterable, List, Tuple, Dict\n", "from typing_extensions import TypedDict\n", + "import re\n", "\n", "from langchain_core.runnables import (\n", " chain as as_runnable,\n", @@ -391,13 +392,20 @@ "\n", "\n", "def _resolve_arg(arg: Union[str, Any], observations: Dict[int, Any]):\n", - " if isinstance(arg, str) and arg.startswith(\"$\"):\n", - " try:\n", - " stripped = arg[1:].replace(\".output\", \"\").strip(\"{}\")\n", - " idx = int(stripped)\n", - " except Exception:\n", - " return str(arg)\n", - " return str(observations[idx])\n", + " # $1 or ${1} -> 1\n", + " ID_PATTERN = r\"\\$\\{?(\\d+)\\}?\"\n", + "\n", + " def replace_match(match):\n", + " # If the string is ${123}, match.group(0) is ${123}, and match.group(1) is 123.\n", + "\n", + " # Return the match group, in this case the index, from the string. This is the index\n", + " # number we get back.\n", + " idx = int(match.group(1))\n", + " return str(observations.get(idx, match.group(0)))\n", + "\n", + " # For dependencies on other tasks\n", + " if isinstance(arg, str):\n", + " return re.sub(ID_PATTERN, replace_match, arg)\n", " elif isinstance(arg, list):\n", " return [_resolve_arg(a, observations) for a in arg]\n", " else:\n", @@ -440,6 +448,7 @@ " # adjust to do a proper topological sort (not-stream)\n", " # or use a more complicated data structure\n", " tasks = scheduler_input[\"tasks\"]\n", + " args_for_tasks = {}\n", " messages = scheduler_input[\"messages\"]\n", " # If we are re-planning, we may have calls that depend on previous\n", " # plans. Start with those.\n", @@ -456,6 +465,7 @@ " task_names[task[\"idx\"]] = (\n", " task[\"tool\"] if isinstance(task[\"tool\"], str) else task[\"tool\"].name\n", " )\n", + " args_for_tasks[task[\"idx\"]] = (task[\"args\"])\n", " if (\n", " # Depends on other tasks\n", " deps\n", @@ -477,12 +487,12 @@ " wait(futures)\n", " # Convert observations to new tool messages to add to the state\n", " new_observations = {\n", - " k: (task_names[k], observations[k])\n", + " k: (task_names[k], args_for_tasks[k], observations[k])\n", " for k in sorted(observations.keys() - originals)\n", " }\n", " tool_messages = [\n", - " FunctionMessage(name=name, content=str(obs), additional_kwargs={\"idx\": k})\n", - " for k, (name, obs) in new_observations.items()\n", + " FunctionMessage(name=name, content=str(obs), additional_kwargs={\"idx\": k, 'args':task_args})\n", + " for k, (name, task_args, obs) in new_observations.items()\n", " ]\n", " return tool_messages" ] @@ -501,7 +511,11 @@ "def plan_and_schedule(messages: List[BaseMessage], config):\n", " tasks = planner.stream(messages, config)\n", " # Begin executing the planner immediately\n", - " tasks = itertools.chain([next(tasks)], tasks)\n", + " try:\n", + " tasks = itertools.chain([next(tasks)], tasks)\n", + " except StopIteration:\n", + " # Handle the case where tasks is empty.\n", + " tasks = iter([])\n", " scheduled_tasks = schedule_tasks.invoke(\n", " {\n", " \"messages\": messages,\n",