Fix Baby AGI notebooks (#2882)

- fix broken notebook cell in
ae485b623d
- Python Black formatting
fix_agent_callbacks
ecneladis 1 year ago committed by GitHub
parent 3c7204d604
commit 1a44b71ddf
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -38,7 +38,7 @@
"from langchain.llms import BaseLLM\n",
"from langchain.vectorstores.base import VectorStore\n",
"from pydantic import BaseModel, Field\n",
"from langchain.chains.base import Chain\n"
"from langchain.chains.base import Chain"
]
},
{
@ -73,6 +73,7 @@
"embeddings_model = OpenAIEmbeddings()\n",
"# Initialize the vectorstore as empty\n",
"import faiss\n",
"\n",
"embedding_size = 1536\n",
"index = faiss.IndexFlatL2(embedding_size)\n",
"vectorstore = FAISS(embeddings_model.embed_query, index, InMemoryDocstore({}), {})"
@ -116,7 +117,12 @@
" )\n",
" prompt = PromptTemplate(\n",
" template=task_creation_template,\n",
" input_variables=[\"result\", \"task_description\", \"incomplete_tasks\", \"objective\"],\n",
" input_variables=[\n",
" \"result\",\n",
" \"task_description\",\n",
" \"incomplete_tasks\",\n",
" \"objective\",\n",
" ],\n",
" )\n",
" return cls(prompt=prompt, llm=llm, verbose=verbose)"
]
@ -147,7 +153,7 @@
" template=task_prioritization_template,\n",
" input_variables=[\"task_names\", \"next_task_id\", \"objective\"],\n",
" )\n",
" return cls(prompt=prompt, llm=llm, verbose=verbose)\n"
" return cls(prompt=prompt, llm=llm, verbose=verbose)"
]
},
{
@ -173,7 +179,7 @@
" template=execution_template,\n",
" input_variables=[\"objective\", \"context\", \"task\"],\n",
" )\n",
" return cls(prompt=prompt, llm=llm, verbose=verbose)\n"
" return cls(prompt=prompt, llm=llm, verbose=verbose)"
]
},
{
@ -193,11 +199,22 @@
"metadata": {},
"outputs": [],
"source": [
"def get_next_task(task_creation_chain: LLMChain, result: Dict, task_description: str, task_list: List[str], objective: str) -> List[Dict]:\n",
"def get_next_task(\n",
" task_creation_chain: LLMChain,\n",
" result: Dict,\n",
" task_description: str,\n",
" task_list: List[str],\n",
" objective: str,\n",
") -> List[Dict]:\n",
" \"\"\"Get the next task.\"\"\"\n",
" incomplete_tasks = \", \".join(task_list)\n",
" response = task_creation_chain.run(result=result, task_description=task_description, incomplete_tasks=incomplete_tasks, objective=objective)\n",
" new_tasks = response.split('\\n')\n",
" response = task_creation_chain.run(\n",
" result=result,\n",
" task_description=task_description,\n",
" incomplete_tasks=incomplete_tasks,\n",
" objective=objective,\n",
" )\n",
" new_tasks = response.split(\"\\n\")\n",
" return [{\"task_name\": task_name} for task_name in new_tasks if task_name.strip()]"
]
},
@ -208,12 +225,19 @@
"metadata": {},
"outputs": [],
"source": [
"def prioritize_tasks(task_prioritization_chain: LLMChain, this_task_id: int, task_list: List[Dict], objective: str) -> List[Dict]:\n",
"def prioritize_tasks(\n",
" task_prioritization_chain: LLMChain,\n",
" this_task_id: int,\n",
" task_list: List[Dict],\n",
" objective: str,\n",
") -> List[Dict]:\n",
" \"\"\"Prioritize tasks.\"\"\"\n",
" task_names = [t[\"task_name\"] for t in task_list]\n",
" next_task_id = int(this_task_id) + 1\n",
" response = task_prioritization_chain.run(task_names=task_names, next_task_id=next_task_id, objective=objective)\n",
" new_tasks = response.split('\\n')\n",
" response = task_prioritization_chain.run(\n",
" task_names=task_names, next_task_id=next_task_id, objective=objective\n",
" )\n",
" new_tasks = response.split(\"\\n\")\n",
" prioritized_task_list = []\n",
" for task_string in new_tasks:\n",
" if not task_string.strip():\n",
@ -239,9 +263,12 @@
" if not results:\n",
" return []\n",
" sorted_results, _ = zip(*sorted(results, key=lambda x: x[1], reverse=True))\n",
" return [str(item.metadata['task']) for item in sorted_results]\n",
" return [str(item.metadata[\"task\"]) for item in sorted_results]\n",
"\n",
"def execute_task(vectorstore, execution_chain: LLMChain, objective: str, task: str, k: int = 5) -> str:\n",
"\n",
"def execute_task(\n",
" vectorstore, execution_chain: LLMChain, objective: str, task: str, k: int = 5\n",
") -> str:\n",
" \"\"\"Execute a task.\"\"\"\n",
" context = _get_top_tasks(vectorstore, query=objective, k=k)\n",
" return execution_chain.run(objective=objective, context=context, task=task)"
@ -254,7 +281,6 @@
"metadata": {},
"outputs": [],
"source": [
"\n",
"class BabyAGI(Chain, BaseModel):\n",
" \"\"\"Controller model for the BabyAGI agent.\"\"\"\n",
"\n",
@ -265,9 +291,10 @@
" task_id_counter: int = Field(1)\n",
" vectorstore: VectorStore = Field(init=False)\n",
" max_iterations: Optional[int] = None\n",
" \n",
"\n",
" class Config:\n",
" \"\"\"Configuration for this pydantic object.\"\"\"\n",
"\n",
" arbitrary_types_allowed = True\n",
"\n",
" def add_task(self, task: Dict):\n",
@ -285,18 +312,18 @@
" def print_task_result(self, result: str):\n",
" print(\"\\033[93m\\033[1m\" + \"\\n*****TASK RESULT*****\\n\" + \"\\033[0m\\033[0m\")\n",
" print(result)\n",
" \n",
"\n",
" @property\n",
" def input_keys(self) -> List[str]:\n",
" return [\"objective\"]\n",
" \n",
"\n",
" @property\n",
" def output_keys(self) -> List[str]:\n",
" return []\n",
"\n",
" def _call(self, inputs: Dict[str, Any]) -> Dict[str, Any]:\n",
" \"\"\"Run the agent.\"\"\"\n",
" objective = inputs['objective']\n",
" objective = inputs[\"objective\"]\n",
" first_task = inputs.get(\"first_task\", \"Make a todo list\")\n",
" self.add_task({\"task_id\": 1, \"task_name\": first_task})\n",
" num_iters = 0\n",
@ -325,7 +352,11 @@
"\n",
" # Step 4: Create new tasks and reprioritize task list\n",
" new_tasks = get_next_task(\n",
" self.task_creation_chain, result, task[\"task_name\"], [t[\"task_name\"] for t in self.task_list], objective\n",
" self.task_creation_chain,\n",
" result,\n",
" task[\"task_name\"],\n",
" [t[\"task_name\"] for t in self.task_list],\n",
" objective,\n",
" )\n",
" for new_task in new_tasks:\n",
" self.task_id_counter += 1\n",
@ -333,27 +364,26 @@
" self.add_task(new_task)\n",
" self.task_list = deque(\n",
" prioritize_tasks(\n",
" self.task_prioritization_chain, this_task_id, list(self.task_list), objective\n",
" self.task_prioritization_chain,\n",
" this_task_id,\n",
" list(self.task_list),\n",
" objective,\n",
" )\n",
" )\n",
" num_iters += 1\n",
" if self.max_iterations is not None and num_iters == self.max_iterations:\n",
" print(\"\\033[91m\\033[1m\" + \"\\n*****TASK ENDING*****\\n\" + \"\\033[0m\\033[0m\")\n",
" print(\n",
" \"\\033[91m\\033[1m\" + \"\\n*****TASK ENDING*****\\n\" + \"\\033[0m\\033[0m\"\n",
" )\n",
" break\n",
" return {}\n",
"\n",
" @classmethod\n",
" def from_llm(\n",
" cls,\n",
" llm: BaseLLM,\n",
" vectorstore: VectorStore,\n",
" verbose: bool = False,\n",
" **kwargs\n",
" cls, llm: BaseLLM, vectorstore: VectorStore, verbose: bool = False, **kwargs\n",
" ) -> \"BabyAGI\":\n",
" \"\"\"Initialize the BabyAGI Controller.\"\"\"\n",
" task_creation_chain = TaskCreationChain.from_llm(\n",
" llm, verbose=verbose\n",
" )\n",
" task_creation_chain = TaskCreationChain.from_llm(llm, verbose=verbose)\n",
" task_prioritization_chain = TaskPrioritizationChain.from_llm(\n",
" llm, verbose=verbose\n",
" )\n",
@ -363,7 +393,7 @@
" task_prioritization_chain=task_prioritization_chain,\n",
" execution_chain=execution_chain,\n",
" vectorstore=vectorstore,\n",
" **kwargs\n",
" **kwargs,\n",
" )"
]
},
@ -405,14 +435,11 @@
"outputs": [],
"source": [
"# Logging of LLMChains\n",
"verbose=False\n",
"verbose = False\n",
"# If None, will keep on going forever\n",
"max_iterations: Optional[int] = 3\n",
"baby_agi = BabyAGI.from_llm(\n",
" llm=llm,\n",
" vectorstore=vectorstore,\n",
" verbose=verbose,\n",
" max_iterations=max_iterations\n",
" llm=llm, vectorstore=vectorstore, verbose=verbose, max_iterations=max_iterations\n",
")"
]
},

@ -34,7 +34,7 @@
"from langchain.llms import BaseLLM\n",
"from langchain.vectorstores.base import VectorStore\n",
"from pydantic import BaseModel, Field\n",
"from langchain.chains.base import Chain\n"
"from langchain.chains.base import Chain"
]
},
{
@ -54,7 +54,9 @@
"metadata": {},
"outputs": [],
"source": [
"%pip install faiss-cpu > /dev/null%pip install google-search-results > /dev/nullfrom langchain.vectorstores import FAISS\n",
"%pip install faiss-cpu > /dev/null\n",
"%pip install google-search-results > /dev/null\n",
"from langchain.vectorstores import FAISS\n",
"from langchain.docstore import InMemoryDocstore"
]
},
@ -69,6 +71,7 @@
"embeddings_model = OpenAIEmbeddings()\n",
"# Initialize the vectorstore as empty\n",
"import faiss\n",
"\n",
"embedding_size = 1536\n",
"index = faiss.IndexFlatL2(embedding_size)\n",
"vectorstore = FAISS(embeddings_model.embed_query, index, InMemoryDocstore({}), {})"
@ -115,7 +118,12 @@
" )\n",
" prompt = PromptTemplate(\n",
" template=task_creation_template,\n",
" input_variables=[\"result\", \"task_description\", \"incomplete_tasks\", \"objective\"],\n",
" input_variables=[\n",
" \"result\",\n",
" \"task_description\",\n",
" \"incomplete_tasks\",\n",
" \"objective\",\n",
" ],\n",
" )\n",
" return cls(prompt=prompt, llm=llm, verbose=verbose)"
]
@ -146,7 +154,7 @@
" template=task_prioritization_template,\n",
" input_variables=[\"task_names\", \"next_task_id\", \"objective\"],\n",
" )\n",
" return cls(prompt=prompt, llm=llm, verbose=verbose)\n"
" return cls(prompt=prompt, llm=llm, verbose=verbose)"
]
},
{
@ -158,20 +166,23 @@
"source": [
"from langchain.agents import ZeroShotAgent, Tool, AgentExecutor\n",
"from langchain import OpenAI, SerpAPIWrapper, LLMChain\n",
"todo_prompt = PromptTemplate.from_template(\"You are a planner who is an expert at coming up with a todo list for a given objective. Come up with a todo list for this objective: {objective}\")\n",
"\n",
"todo_prompt = PromptTemplate.from_template(\n",
" \"You are a planner who is an expert at coming up with a todo list for a given objective. Come up with a todo list for this objective: {objective}\"\n",
")\n",
"todo_chain = LLMChain(llm=OpenAI(temperature=0), prompt=todo_prompt)\n",
"search = SerpAPIWrapper()\n",
"tools = [\n",
" Tool(\n",
" name = \"Search\",\n",
" name=\"Search\",\n",
" func=search.run,\n",
" description=\"useful for when you need to answer questions about current events\"\n",
" description=\"useful for when you need to answer questions about current events\",\n",
" ),\n",
" Tool(\n",
" name = \"TODO\",\n",
" name=\"TODO\",\n",
" func=todo_chain.run,\n",
" description=\"useful for when you need to come up with todo lists. Input: an objective to create a todo list for. Output: a todo list for that objective. Please be very clear what the objective is!\"\n",
" )\n",
" description=\"useful for when you need to come up with todo lists. Input: an objective to create a todo list for. Output: a todo list for that objective. Please be very clear what the objective is!\",\n",
" ),\n",
"]\n",
"\n",
"\n",
@ -179,10 +190,10 @@
"suffix = \"\"\"Question: {task}\n",
"{agent_scratchpad}\"\"\"\n",
"prompt = ZeroShotAgent.create_prompt(\n",
" tools, \n",
" prefix=prefix, \n",
" suffix=suffix, \n",
" input_variables=[\"objective\", \"task\", \"context\",\"agent_scratchpad\"]\n",
" tools,\n",
" prefix=prefix,\n",
" suffix=suffix,\n",
" input_variables=[\"objective\", \"task\", \"context\", \"agent_scratchpad\"],\n",
")"
]
},
@ -203,11 +214,22 @@
"metadata": {},
"outputs": [],
"source": [
"def get_next_task(task_creation_chain: LLMChain, result: Dict, task_description: str, task_list: List[str], objective: str) -> List[Dict]:\n",
"def get_next_task(\n",
" task_creation_chain: LLMChain,\n",
" result: Dict,\n",
" task_description: str,\n",
" task_list: List[str],\n",
" objective: str,\n",
") -> List[Dict]:\n",
" \"\"\"Get the next task.\"\"\"\n",
" incomplete_tasks = \", \".join(task_list)\n",
" response = task_creation_chain.run(result=result, task_description=task_description, incomplete_tasks=incomplete_tasks, objective=objective)\n",
" new_tasks = response.split('\\n')\n",
" response = task_creation_chain.run(\n",
" result=result,\n",
" task_description=task_description,\n",
" incomplete_tasks=incomplete_tasks,\n",
" objective=objective,\n",
" )\n",
" new_tasks = response.split(\"\\n\")\n",
" return [{\"task_name\": task_name} for task_name in new_tasks if task_name.strip()]"
]
},
@ -218,12 +240,19 @@
"metadata": {},
"outputs": [],
"source": [
"def prioritize_tasks(task_prioritization_chain: LLMChain, this_task_id: int, task_list: List[Dict], objective: str) -> List[Dict]:\n",
"def prioritize_tasks(\n",
" task_prioritization_chain: LLMChain,\n",
" this_task_id: int,\n",
" task_list: List[Dict],\n",
" objective: str,\n",
") -> List[Dict]:\n",
" \"\"\"Prioritize tasks.\"\"\"\n",
" task_names = [t[\"task_name\"] for t in task_list]\n",
" next_task_id = int(this_task_id) + 1\n",
" response = task_prioritization_chain.run(task_names=task_names, next_task_id=next_task_id, objective=objective)\n",
" new_tasks = response.split('\\n')\n",
" response = task_prioritization_chain.run(\n",
" task_names=task_names, next_task_id=next_task_id, objective=objective\n",
" )\n",
" new_tasks = response.split(\"\\n\")\n",
" prioritized_task_list = []\n",
" for task_string in new_tasks:\n",
" if not task_string.strip():\n",
@ -249,9 +278,12 @@
" if not results:\n",
" return []\n",
" sorted_results, _ = zip(*sorted(results, key=lambda x: x[1], reverse=True))\n",
" return [str(item.metadata['task']) for item in sorted_results]\n",
" return [str(item.metadata[\"task\"]) for item in sorted_results]\n",
"\n",
"def execute_task(vectorstore, execution_chain: LLMChain, objective: str, task: str, k: int = 5) -> str:\n",
"\n",
"def execute_task(\n",
" vectorstore, execution_chain: LLMChain, objective: str, task: str, k: int = 5\n",
") -> str:\n",
" \"\"\"Execute a task.\"\"\"\n",
" context = _get_top_tasks(vectorstore, query=objective, k=k)\n",
" return execution_chain.run(objective=objective, context=context, task=task)"
@ -264,7 +296,6 @@
"metadata": {},
"outputs": [],
"source": [
"\n",
"class BabyAGI(Chain, BaseModel):\n",
" \"\"\"Controller model for the BabyAGI agent.\"\"\"\n",
"\n",
@ -275,9 +306,10 @@
" task_id_counter: int = Field(1)\n",
" vectorstore: VectorStore = Field(init=False)\n",
" max_iterations: Optional[int] = None\n",
" \n",
"\n",
" class Config:\n",
" \"\"\"Configuration for this pydantic object.\"\"\"\n",
"\n",
" arbitrary_types_allowed = True\n",
"\n",
" def add_task(self, task: Dict):\n",
@ -295,18 +327,18 @@
" def print_task_result(self, result: str):\n",
" print(\"\\033[93m\\033[1m\" + \"\\n*****TASK RESULT*****\\n\" + \"\\033[0m\\033[0m\")\n",
" print(result)\n",
" \n",
"\n",
" @property\n",
" def input_keys(self) -> List[str]:\n",
" return [\"objective\"]\n",
" \n",
"\n",
" @property\n",
" def output_keys(self) -> List[str]:\n",
" return []\n",
"\n",
" def _call(self, inputs: Dict[str, Any]) -> Dict[str, Any]:\n",
" \"\"\"Run the agent.\"\"\"\n",
" objective = inputs['objective']\n",
" objective = inputs[\"objective\"]\n",
" first_task = inputs.get(\"first_task\", \"Make a todo list\")\n",
" self.add_task({\"task_id\": 1, \"task_name\": first_task})\n",
" num_iters = 0\n",
@ -335,7 +367,11 @@
"\n",
" # Step 4: Create new tasks and reprioritize task list\n",
" new_tasks = get_next_task(\n",
" self.task_creation_chain, result, task[\"task_name\"], [t[\"task_name\"] for t in self.task_list], objective\n",
" self.task_creation_chain,\n",
" result,\n",
" task[\"task_name\"],\n",
" [t[\"task_name\"] for t in self.task_list],\n",
" objective,\n",
" )\n",
" for new_task in new_tasks:\n",
" self.task_id_counter += 1\n",
@ -343,40 +379,41 @@
" self.add_task(new_task)\n",
" self.task_list = deque(\n",
" prioritize_tasks(\n",
" self.task_prioritization_chain, this_task_id, list(self.task_list), objective\n",
" self.task_prioritization_chain,\n",
" this_task_id,\n",
" list(self.task_list),\n",
" objective,\n",
" )\n",
" )\n",
" num_iters += 1\n",
" if self.max_iterations is not None and num_iters == self.max_iterations:\n",
" print(\"\\033[91m\\033[1m\" + \"\\n*****TASK ENDING*****\\n\" + \"\\033[0m\\033[0m\")\n",
" print(\n",
" \"\\033[91m\\033[1m\" + \"\\n*****TASK ENDING*****\\n\" + \"\\033[0m\\033[0m\"\n",
" )\n",
" break\n",
" return {}\n",
"\n",
" @classmethod\n",
" def from_llm(\n",
" cls,\n",
" llm: BaseLLM,\n",
" vectorstore: VectorStore,\n",
" verbose: bool = False,\n",
" **kwargs\n",
" cls, llm: BaseLLM, vectorstore: VectorStore, verbose: bool = False, **kwargs\n",
" ) -> \"BabyAGI\":\n",
" \"\"\"Initialize the BabyAGI Controller.\"\"\"\n",
" task_creation_chain = TaskCreationChain.from_llm(\n",
" llm, verbose=verbose\n",
" )\n",
" task_creation_chain = TaskCreationChain.from_llm(llm, verbose=verbose)\n",
" task_prioritization_chain = TaskPrioritizationChain.from_llm(\n",
" llm, verbose=verbose\n",
" )\n",
" llm_chain = LLMChain(llm=llm, prompt=prompt)\n",
" tool_names = [tool.name for tool in tools]\n",
" agent = ZeroShotAgent(llm_chain=llm_chain, allowed_tools=tool_names)\n",
" agent_executor = AgentExecutor.from_agent_and_tools(agent=agent, tools=tools, verbose=True)\n",
" agent_executor = AgentExecutor.from_agent_and_tools(\n",
" agent=agent, tools=tools, verbose=True\n",
" )\n",
" return cls(\n",
" task_creation_chain=task_creation_chain,\n",
" task_prioritization_chain=task_prioritization_chain,\n",
" execution_chain=agent_executor,\n",
" vectorstore=vectorstore,\n",
" **kwargs\n",
" **kwargs,\n",
" )"
]
},
@ -418,14 +455,11 @@
"outputs": [],
"source": [
"# Logging of LLMChains\n",
"verbose=False\n",
"verbose = False\n",
"# If None, will keep on going forever\n",
"max_iterations: Optional[int] = 3\n",
"baby_agi = BabyAGI.from_llm(\n",
" llm=llm,\n",
" vectorstore=vectorstore,\n",
" verbose=verbose,\n",
" max_iterations=max_iterations\n",
" llm=llm, vectorstore=vectorstore, verbose=verbose, max_iterations=max_iterations\n",
")"
]
},

Loading…
Cancel
Save