├── .gitignore ├── README.md └── agent.py /.gitignore: -------------------------------------------------------------------------------- 1 | venv 2 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ## Overview 2 | 3 | The Agent creates a Todo list to complete tasks that are required by the user and even searches Google for you. 4 | This is a BabyAGI + SerpAPI tool that uses Langchain's framework. The site is hosted by Streamlit. 5 | 6 | --- 7 | 8 | ## To run 9 | 10 | Make sure to fill in the API keys for OpenAI and SerpAPI. 11 | 12 | ``` 13 | python -m venv venv 14 | source venv/bin/activate 15 | pip install langchain 16 | pip install openai 17 | pip install google-search-results 18 | pip install faiss-cpu > /dev/null 19 | pip install tiktoken 20 | pip install streamlit 21 | streamlit run agent.py 22 | 23 | ``` 24 | If streamlit is not working, checkout their [installation page](https://docs.streamlit.io/library/get-started/installation) 25 | 26 | --- 27 | 28 | ### Develop Log 29 | 30 | 23-04-18 31 | 32 | 1. BabyAGI + serpapi Agent 33 | 34 | 2. Frontend UI on streamlit 35 | - User input 36 | - Download results as one .txt file 37 | 38 | --- 39 | 40 | ### TODO 41 | 42 | - [ ] Need to revise the prompts for more precise tasks 43 | 44 | --- 45 | -------------------------------------------------------------------------------- /agent.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ["OPENAI_API_KEY"] = "" 3 | os.environ["SERPAPI_API_KEY"] = "" 4 | import streamlit as st 5 | from collections import deque 6 | from typing import Dict, List, Optional, Any 7 | 8 | from langchain import LLMChain, OpenAI, PromptTemplate 9 | from langchain.embeddings import OpenAIEmbeddings 10 | from langchain.llms import BaseLLM 11 | from langchain.vectorstores.base import VectorStore 12 | from pydantic import BaseModel, Field 13 | from langchain.chains.base import Chain 14 | from langchain.vectorstores import FAISS 15 | from langchain.docstore import InMemoryDocstore 16 | 17 | st.title('Personal AI Assistant') 18 | 19 | # Define your embedding model 20 | embeddings_model = OpenAIEmbeddings() 21 | # Initialize the vectorstore as empty 22 | import faiss 23 | 24 | embedding_size = 1536 25 | index = faiss.IndexFlatL2(embedding_size) 26 | vectorstore = FAISS(embeddings_model.embed_query, index, InMemoryDocstore({}), {}) 27 | 28 | class TaskCreationChain(LLMChain): 29 | """Chain to generates tasks.""" 30 | 31 | @classmethod 32 | def from_llm(cls, llm: BaseLLM, verbose: bool = True) -> LLMChain: 33 | """Get the response parser.""" 34 | task_creation_template = ( 35 | "You are an task creation AI that uses the result of an execution agent" 36 | " to create new tasks with the following objective: {objective}," 37 | " The last completed task has the result: {result}." 38 | " This result was based on this task description: {task_description}." 39 | " These are incomplete tasks: {incomplete_tasks}." 40 | " Based on the result, create new tasks to be completed" 41 | " by the AI system that do not overlap with incomplete tasks." 42 | " Return the tasks as an array." 43 | ) 44 | prompt = PromptTemplate( 45 | template=task_creation_template, 46 | input_variables=[ 47 | "result", 48 | "task_description", 49 | "incomplete_tasks", 50 | "objective", 51 | ], 52 | ) 53 | return cls(prompt=prompt, llm=llm, verbose=verbose) 54 | 55 | class TaskPrioritizationChain(LLMChain): 56 | """Chain to prioritize tasks.""" 57 | 58 | @classmethod 59 | def from_llm(cls, llm: BaseLLM, verbose: bool = True) -> LLMChain: 60 | """Get the response parser.""" 61 | task_prioritization_template = ( 62 | "You are an task prioritization AI tasked with cleaning the formatting of and reprioritizing" 63 | " the following tasks: {task_names}." 64 | " Consider the ultimate objective of your team: {objective}." 65 | " Do not remove any tasks. Return the result as a numbered list, like:" 66 | " #. First task" 67 | " #. Second task" 68 | " Start the task list with number {next_task_id}." 69 | ) 70 | prompt = PromptTemplate( 71 | template=task_prioritization_template, 72 | input_variables=["task_names", "next_task_id", "objective"], 73 | ) 74 | return cls(prompt=prompt, llm=llm, verbose=verbose) 75 | 76 | from langchain.agents import ZeroShotAgent, Tool, AgentExecutor 77 | from langchain import OpenAI, SerpAPIWrapper, LLMChain 78 | 79 | todo_prompt = PromptTemplate.from_template( 80 | "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}" 81 | ) 82 | todo_chain = LLMChain(llm=OpenAI(temperature=0), prompt=todo_prompt) 83 | search = SerpAPIWrapper() 84 | tools = [ 85 | Tool( 86 | name="Search", 87 | func=search.run, 88 | description="useful for when you need to answer questions about current events", 89 | ), 90 | Tool( 91 | name="TODO", 92 | func=todo_chain.run, 93 | 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!", 94 | ), 95 | ] 96 | 97 | 98 | prefix = """You are an AI who performs one task based on the following objective: {objective}. Take into account these previously completed tasks: {context}.""" 99 | suffix = """Question: {task} 100 | {agent_scratchpad}""" 101 | prompt = ZeroShotAgent.create_prompt( 102 | tools, 103 | prefix=prefix, 104 | suffix=suffix, 105 | input_variables=["objective", "task", "context", "agent_scratchpad"], 106 | ) 107 | 108 | def get_next_task( 109 | task_creation_chain: LLMChain, 110 | result: Dict, 111 | task_description: str, 112 | task_list: List[str], 113 | objective: str, 114 | ) -> List[Dict]: 115 | """Get the next task.""" 116 | incomplete_tasks = ", ".join(task_list) 117 | response = task_creation_chain.run( 118 | result=result, 119 | task_description=task_description, 120 | incomplete_tasks=incomplete_tasks, 121 | objective=objective, 122 | ) 123 | new_tasks = response.split("\n") 124 | return [{"task_name": task_name} for task_name in new_tasks if task_name.strip()] 125 | 126 | def prioritize_tasks( 127 | task_prioritization_chain: LLMChain, 128 | this_task_id: int, 129 | task_list: List[Dict], 130 | objective: str, 131 | ) -> List[Dict]: 132 | """Prioritize tasks.""" 133 | task_names = [t["task_name"] for t in task_list] 134 | next_task_id = int(this_task_id) + 1 135 | response = task_prioritization_chain.run( 136 | task_names=task_names, next_task_id=next_task_id, objective=objective 137 | ) 138 | new_tasks = response.split("\n") 139 | prioritized_task_list = [] 140 | for task_string in new_tasks: 141 | if not task_string.strip(): 142 | continue 143 | task_parts = task_string.strip().split(".", 1) 144 | if len(task_parts) == 2: 145 | task_id = task_parts[0].strip() 146 | task_name = task_parts[1].strip() 147 | prioritized_task_list.append({"task_id": task_id, "task_name": task_name}) 148 | return prioritized_task_list 149 | 150 | def _get_top_tasks(vectorstore, query: str, k: int) -> List[str]: 151 | """Get the top k tasks based on the query.""" 152 | results = vectorstore.similarity_search_with_score(query, k=k) 153 | if not results: 154 | return [] 155 | sorted_results, _ = zip(*sorted(results, key=lambda x: x[1], reverse=True)) 156 | return [str(item.metadata["task"]) for item in sorted_results] 157 | 158 | 159 | def execute_task( 160 | vectorstore, execution_chain: LLMChain, objective: str, task: str, k: int = 5 161 | ) -> str: 162 | """Execute a task.""" 163 | context = _get_top_tasks(vectorstore, query=objective, k=k) 164 | return execution_chain.run(objective=objective, context=context, task=task) 165 | 166 | class BabyAGI(Chain, BaseModel): 167 | """Controller model for the BabyAGI agent.""" 168 | 169 | task_list: deque = Field(default_factory=deque) 170 | task_creation_chain: TaskCreationChain = Field(...) 171 | task_prioritization_chain: TaskPrioritizationChain = Field(...) 172 | execution_chain: AgentExecutor = Field(...) 173 | task_id_counter: int = Field(1) 174 | vectorstore: VectorStore = Field(init=False) 175 | max_iterations: Optional[int] = None 176 | 177 | class Config: 178 | """Configuration for this pydantic object.""" 179 | 180 | arbitrary_types_allowed = True 181 | 182 | def add_task(self, task: Dict): 183 | self.task_list.append(task) 184 | 185 | def print_task_list(self): 186 | print("\033[95m\033[1m" + "\n*****TASK LIST*****\n" + "\033[0m\033[0m") 187 | for t in self.task_list: 188 | print(str(t["task_id"]) + ": " + t["task_name"]) 189 | 190 | def print_next_task(self, task: Dict): 191 | print("\033[92m\033[1m" + "\n*****NEXT TASK*****\n" + "\033[0m\033[0m") 192 | print(str(task["task_id"]) + ": " + task["task_name"]) 193 | return (str(task["task_id"]) + ": " + task["task_name"]) 194 | 195 | def print_task_result(self, result: str): 196 | print("\033[93m\033[1m" + "\n*****TASK RESULT*****\n" + "\033[0m\033[0m") 197 | return(result) 198 | 199 | @property 200 | def input_keys(self) -> List[str]: 201 | return ["objective"] 202 | 203 | @property 204 | def output_keys(self) -> List[str]: 205 | return [] 206 | 207 | @st.cache_data 208 | def _call(_self, inputs: Dict[str, Any]) -> Dict[str, Any]: 209 | result_list = [] 210 | """Run the agent.""" 211 | objective = inputs["objective"] 212 | first_task = inputs.get("first_task", "Make a todo list") 213 | _self.add_task({"task_id": 1, "task_name": first_task}) 214 | num_iters = 0 215 | while True: 216 | if _self.task_list: 217 | _self.print_task_list() 218 | 219 | # Step 1: Pull the first task 220 | task = _self.task_list.popleft() 221 | _self.print_next_task(task) 222 | st.write('**Task:** \n') 223 | st.write(_self.print_next_task(task)) 224 | 225 | # Step 2: Execute the task 226 | result = execute_task( 227 | _self.vectorstore, _self.execution_chain, objective, task["task_name"] 228 | ) 229 | this_task_id = int(task["task_id"]) 230 | _self.print_task_result(result) 231 | st.write('**Result:** \n') 232 | st.write(_self.print_task_result(result)) 233 | result_list.append(result) 234 | 235 | # Step 3: Store the result in Pinecone 236 | result_id = f"result_{task['task_id']}" 237 | _self.vectorstore.add_texts( 238 | texts=[result], 239 | metadatas=[{"task": task["task_name"]}], 240 | ids=[result_id], 241 | ) 242 | 243 | # Step 4: Create new tasks and reprioritize task list 244 | new_tasks = get_next_task( 245 | _self.task_creation_chain, 246 | result, 247 | task["task_name"], 248 | [t["task_name"] for t in _self.task_list], 249 | objective, 250 | ) 251 | for new_task in new_tasks: 252 | _self.task_id_counter += 1 253 | new_task.update({"task_id": _self.task_id_counter}) 254 | _self.add_task(new_task) 255 | _self.task_list = deque( 256 | prioritize_tasks( 257 | _self.task_prioritization_chain, 258 | this_task_id, 259 | list(_self.task_list), 260 | objective, 261 | ) 262 | ) 263 | num_iters += 1 264 | if _self.max_iterations is not None and num_iters == _self.max_iterations: 265 | print( 266 | "\033[91m\033[1m" + "\n*****TASK ENDING*****\n" + "\033[0m\033[0m" 267 | ) 268 | st.success('Task Completed!', icon="✅") 269 | break 270 | 271 | # Create a temporary file to hold the text 272 | with open('output.txt', 'w') as f: 273 | for item in result_list: 274 | f.write(item) 275 | f.write("\n\n") 276 | 277 | return {} 278 | 279 | @classmethod 280 | def from_llm( 281 | cls, llm: BaseLLM, vectorstore: VectorStore, verbose: bool = False, **kwargs 282 | ) -> "BabyAGI": 283 | """Initialize the BabyAGI Controller.""" 284 | task_creation_chain = TaskCreationChain.from_llm(llm, verbose=verbose) 285 | task_prioritization_chain = TaskPrioritizationChain.from_llm( 286 | llm, verbose=verbose 287 | ) 288 | llm_chain = LLMChain(llm=llm, prompt=prompt) 289 | tool_names = [tool.name for tool in tools] 290 | agent = ZeroShotAgent(llm_chain=llm_chain, allowed_tools=tool_names) 291 | agent_executor = AgentExecutor.from_agent_and_tools( 292 | agent=agent, tools=tools, verbose=True 293 | ) 294 | return cls( 295 | task_creation_chain=task_creation_chain, 296 | task_prioritization_chain=task_prioritization_chain, 297 | execution_chain=agent_executor, 298 | vectorstore=vectorstore, 299 | **kwargs, 300 | ) 301 | 302 | 303 | def get_text(): 304 | input_text = st.text_input("Type in your prompt below", key="input") 305 | return input_text 306 | 307 | user_input = get_text() 308 | 309 | OBJECTIVE = user_input 310 | llm = OpenAI(temperature=0) 311 | # Logging of LLMChains 312 | verbose = False 313 | # If None, will keep on going forever. Customize the number of loops you want it to go through. 314 | max_iterations: Optional[int] = 2 315 | baby_agi = BabyAGI.from_llm( 316 | llm=llm, vectorstore=vectorstore, verbose=verbose, max_iterations=max_iterations 317 | ) 318 | 319 | if (user_input): 320 | baby_agi({"objective": OBJECTIVE}) 321 | 322 | # Download the file using Streamlit's download_button() function 323 | st.download_button( 324 | label='Download Results', 325 | data=open('output.txt', 'rb').read(), 326 | file_name='output.txt', 327 | mime='text/plain' 328 | ) 329 | 330 | # baby_agi({"objective": OBJECTIVE}) 331 | --------------------------------------------------------------------------------