from langchain.vectorstores.base import VectorStore from pydantic import BaseModel, Field from langchain.chains.base import Chain from collections import deque from typing import Dict, List, Optional, Any from langchain.agents import ZeroShotAgent, AgentExecutor from src.task_creation_chain import TaskCreationChain from src.task_prio_chain import TaskPrioritizationChain import streamlit as st from langchain import LLMChain from langchain.llms import BaseLLM # -----------------helpers def get_next_task( task_creation_chain: LLMChain, result: Dict, task_description: str, task_list: List[str], objective: str, ) -> List[Dict]: """Get the next task.""" incomplete_tasks = ", ".join(task_list) response = task_creation_chain.run( result=result, task_description=task_description, incomplete_tasks=incomplete_tasks, objective=objective, ) new_tasks = response.split("\n") return [{"task_name": task_name} for task_name in new_tasks if task_name.strip()] def prioritize_tasks( task_prioritization_chain: LLMChain, this_task_id: int, task_list: List[Dict], objective: str, ) -> List[Dict]: """Prioritize tasks.""" task_names = [t["task_name"] for t in task_list] next_task_id = int(this_task_id) + 1 response = task_prioritization_chain.run( task_names=task_names, next_task_id=next_task_id, objective=objective ) new_tasks = response.split("\n") prioritized_task_list = [] for task_string in new_tasks: if not task_string.strip(): continue task_parts = task_string.strip().split(".", 1) if len(task_parts) == 2: task_id = task_parts[0].strip() task_name = task_parts[1].strip() prioritized_task_list.append({"task_id": task_id, "task_name": task_name}) return prioritized_task_list def _get_top_tasks(vectorstore, query: str, k: int) -> List[str]: """Get the top k tasks based on the query.""" results = vectorstore.similarity_search_with_score(query, k=k) if not results: return [] sorted_results, _ = zip(*sorted(results, key=lambda x: x[1], reverse=True)) return [str(item.metadata["task"]) for item in sorted_results] def execute_task( vectorstore, execution_chain: LLMChain, objective: str, task: str, k: int = 5 ) -> str: """Execute a task.""" context = _get_top_tasks(vectorstore, query=objective, k=k) return execution_chain.run(objective=objective, context=context, task=task) # ---------------Class------------- class BabyAGI(Chain, BaseModel): """Controller model for the BabyAGI agent.""" task_list: deque = Field(default_factory=deque) task_creation_chain: TaskCreationChain = Field(...) task_prioritization_chain: TaskPrioritizationChain = Field(...) execution_chain: AgentExecutor = Field(...) task_id_counter: int = Field(1) vectorstore: VectorStore = Field(init=False) max_iterations: Optional[int] = None class Config: """Configuration for this pydantic object.""" arbitrary_types_allowed = True def add_task(self, task: Dict): self.task_list.append(task) def print_task_list(self): print("\033[95m\033[1m" + "\n*****TASK LIST*****\n" + "\033[0m\033[0m") if len(self.task_list) > 1: st.write('**Task List:** \n') for t in self.task_list: print(str(t["task_id"]) + ": " + t["task_name"]) if len(self.task_list) > 1: st.write(str(t["task_id"]) + ": " + t["task_name"]) def print_next_task(self, task: Dict): print("\033[92m\033[1m" + "\n*****NEXT TASK*****\n" + "\033[0m\033[0m") print(str(task["task_name"])) return (str(task["task_name"])) def print_task_result(self, result: str): print("\033[93m\033[1m" + "\n*****TASK RESULT*****\n" + "\033[0m\033[0m") return(result) @property def input_keys(self) -> List[str]: return ["objective"] @property def output_keys(self) -> List[str]: return [] def _call(_self, inputs: Dict[str, Any]) -> Dict[str, Any]: result_list = [] """Run the agent.""" objective = inputs["objective"] first_task = inputs.get("first_task", f"Make a todo list to accomplish the objective: {objective}") _self.add_task({"task_id": 1, "task_name": first_task}) num_iters = 0 while True: if _self.task_list: _self.print_task_list() # Step 1: Pull the first task task = _self.task_list.popleft() _self.print_next_task(task) st.write('**Next Task:** \n') st.write(_self.print_next_task(task)) # Step 2: Execute the task result = execute_task( _self.vectorstore, _self.execution_chain, objective, task["task_name"] ) this_task_id = int(task["task_id"]) _self.print_task_result(result) st.write('**Result from Task:** \n') st.write(_self.print_task_result(result)) result_list.append(result) # Step 3: Store the result in Pinecone result_id = f"result_{task['task_id']}" _self.vectorstore.add_texts( texts=[result], metadatas=[{"task": task["task_name"]}], ids=[result_id], ) # Step 4: Create new tasks and reprioritize task list new_tasks = get_next_task( _self.task_creation_chain, result, task["task_name"], [t["task_name"] for t in _self.task_list], objective, ) for new_task in new_tasks: _self.task_id_counter += 1 new_task.update({"task_id": _self.task_id_counter}) _self.add_task(new_task) _self.task_list = deque( prioritize_tasks( _self.task_prioritization_chain, this_task_id, list(_self.task_list), objective, ) ) num_iters += 1 if _self.max_iterations is not None and num_iters == _self.max_iterations: print( "\033[91m\033[1m" + "\n*****TASK ENDING*****\n" + "\033[0m\033[0m" ) st.success('Task Completed!', icon="✅") break # Create a temporary file to hold the text with open('output.txt', 'w') as f: for item in result_list: f.write(item) f.write("\n\n") return {} @classmethod def from_llm( cls, prompt: str, tools: list, llm: BaseLLM, vectorstore: VectorStore, verbose: bool = False, **kwargs ) -> "BabyAGI": """Initialize the BabyAGI Controller.""" task_creation_chain = TaskCreationChain.from_llm(llm, verbose=verbose) task_prioritization_chain = TaskPrioritizationChain.from_llm( llm, verbose=verbose ) llm_chain = LLMChain(llm=llm, prompt=prompt) tool_names = [tool.name for tool in tools] agent = ZeroShotAgent(llm_chain=llm_chain, allowed_tools=tool_names) agent_executor = AgentExecutor.from_agent_and_tools( agent=agent, tools=tools, verbose=True ) return cls( task_creation_chain=task_creation_chain, task_prioritization_chain=task_prioritization_chain, execution_chain=agent_executor, vectorstore=vectorstore, **kwargs, )