Spaces:
Runtime error
Runtime error
| import logging | |
| import os | |
| import faiss | |
| from langchain import LLMChain, OpenAI, PromptTemplate | |
| from langchain.agents import AgentExecutor, Tool, ZeroShotAgent | |
| from langchain.docstore import InMemoryDocstore | |
| from langchain.embeddings import OpenAIEmbeddings | |
| from langchain.vectorstores import FAISS | |
| from langchain_experimental.autonomous_agents import BabyAGI | |
| from pydantic import ValidationError | |
| logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') | |
| # ---------- Boss Node ---------- | |
| class Boss: | |
| """ | |
| The Bose class is responsible for creating and executing tasks using the BabyAGI model. | |
| It takes a language model (llm), a vectorstore for memory, an agent_executor for task execution, and a maximum number of iterations for the BabyAGI model. | |
| # Setup | |
| api_key = "YOUR_OPENAI_API_KEY" # Replace with your OpenAI API Key. | |
| os.environ["OPENAI_API_KEY"] = api_key | |
| # Objective for the Boss | |
| objective = "Analyze website user behavior patterns over the past month." | |
| # Create a Bose instance | |
| boss = Bose( | |
| objective=objective, | |
| boss_system_prompt="You are the main controller of a data analysis swarm...", | |
| api_key=api_key, | |
| worker_node=WorkerNode | |
| ) | |
| # Run the Bose to process the objective | |
| boss.run() | |
| """ | |
| def __init__( | |
| self, | |
| objective: str, | |
| api_key=None, | |
| max_iterations=5, | |
| human_in_the_loop=None, | |
| boss_system_prompt="You are a boss planner in a swarm...", | |
| llm_class=OpenAI, | |
| worker_node=None, | |
| verbose=False | |
| ): | |
| # Store parameters | |
| self.api_key = api_key or os.getenv("OPENAI_API_KEY") | |
| self.objective = objective | |
| self.max_iterations = max_iterations | |
| self.boss_system_prompt = boss_system_prompt | |
| self.llm_class = llm_class | |
| self.verbose = verbose | |
| # Initialization methods | |
| self.llm = self._initialize_llm() | |
| self.vectorstore = self._initialize_vectorstore() | |
| self.task = self._create_task(self.objective) | |
| self.agent_executor = self._initialize_agent_executor(worker_node) | |
| self.baby_agi = self._initialize_baby_agi(human_in_the_loop) | |
| def _initialize_llm(self): | |
| """ | |
| Init LLM | |
| Params: | |
| llm_class(class): The Language model class. Default is OpenAI. | |
| temperature (float): The Temperature for the language model. Default is 0.5 | |
| """ | |
| try: | |
| return self.llm_class(openai_api_key=self.api_key, temperature=0.5) | |
| except Exception as e: | |
| logging.error(f"Failed to initialize language model: {e}") | |
| raise e | |
| def _initialize_vectorstore(self): | |
| try: | |
| embeddings_model = OpenAIEmbeddings(openai_api_key=self.api_key) | |
| embedding_size = 8192 | |
| index = faiss.IndexFlatL2(embedding_size) | |
| return FAISS( | |
| embeddings_model.embed_query, | |
| index, | |
| InMemoryDocstore({}), {} | |
| ) | |
| except Exception as e: | |
| logging.error(f"Failed to initialize vector store: {e}") | |
| raise e | |
| def _initialize_agent_executor(self, worker_node): | |
| todo_prompt = PromptTemplate.from_template(self.boss_system_prompt) | |
| todo_chain = LLMChain(llm=self.llm, prompt=todo_prompt) | |
| tools = [ | |
| Tool( | |
| name="Goal Decomposition Tool", | |
| func=todo_chain.run, | |
| description="Use Case: Decompose ambitious goals into as many explicit and well defined tasks for an AI agent to follow. Rules and Regulations, don't use this tool too often only in the beginning when the user grants you a mission." | |
| ), | |
| Tool(name="Swarm Worker Agent", func=worker_node, description="Use Case: When you want to delegate and assign the decomposed goal sub tasks to a worker agent in your swarm, Rules and Regulations, Provide a task specification sheet to the worker agent. It can use the browser, process csvs and generate content") | |
| ] | |
| suffix = """Question: {task}\n{agent_scratchpad}""" | |
| prefix = """You are a Boss in a swarm who performs one task based on the following objective: {objective}. Take into account these previously completed tasks: {context}.\n """ | |
| prompt = ZeroShotAgent.create_prompt( | |
| tools, | |
| prefix=prefix, | |
| suffix=suffix, | |
| input_variables=["objective", "task", "context", "agent_scratchpad"], | |
| ) | |
| llm_chain = LLMChain(llm=self.llm, prompt=prompt) | |
| agent = ZeroShotAgent(llm_chain=llm_chain, allowed_tools=tools) | |
| return AgentExecutor.from_agent_and_tools(agent=agent, tools=tools, verbose=self.verbose) | |
| def _initialize_baby_agi(self, human_in_the_loop): | |
| try: | |
| return BabyAGI.from_llm( | |
| llm=self.llm, | |
| vectorstore=self.vectorstore, | |
| task_execution_chain=self.agent_executor, | |
| max_iterations=self.max_iterations, | |
| human_in_the_loop=human_in_the_loop | |
| ) | |
| except ValidationError as e: | |
| logging.error(f"Validation Error while initializing BabyAGI: {e}") | |
| raise | |
| except Exception as e: | |
| logging.error(f"Unexpected Error while initializing BabyAGI: {e}") | |
| raise | |
| def _create_task(self, objective): | |
| if not objective: | |
| logging.error("Objective cannot be empty.") | |
| raise ValueError("Objective cannot be empty.") | |
| return {"objective": objective} | |
| def run(self): | |
| if not self.task: | |
| logging.error("Task cannot be empty.") | |
| raise ValueError("Task cannot be empty.") | |
| try: | |
| self.baby_agi(self.task) | |
| except Exception as e: | |
| logging.error(f"Error while executing task: {e}") | |
| raise | |