Merge pull request yoheinakajima#25 from francip/main

Add support for GPT-4

README.md

@@ -33,7 +33,10 @@ To use the script, you will need to follow these steps:
 3. Set your OpenAI and Pinecone API keys in the OPENAI_API_KEY and PINECONE_API_KEY variables.
 4. Set the Pinecone environment in the PINECONE_ENVIRONMENT variable.
 5. Set the name of the table where the task results will be stored in the TABLE_NAME variable.
-6. Set the objective of the task management system in the OBJECTIVE variable.
+6. Set the objective of the task management system in the OBJECTIVE variable. Alternatively you can pass it to the script as a quote argument.
+```
+./babyagi.py ["<objective>"]
+```
 7. Set the first task of the system in the FIRST_TASK variable.
 8. Run the script.
 

babyagi.py

@@ -1,17 +1,26 @@
+#!/usr/bin/env python3
+import os
 import openai
 import pinecone
 import time
+import sys
 from collections import deque
 from typing import Dict, List
 from dotenv import load_dotenv
 import os
 
+#Set Variables
 load_dotenv()
 
 # Set API Keys
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "")
 assert OPENAI_API_KEY, "OPENAI_API_KEY environment variable is missing from .env from .env"
 
+# Use GPT-3 model
+USE_GPT4 = False
+if USE_GPT4:
+    print("\033[91m\033[1m"+"\n*****USING GPT-4. POTENTIALLY EXPENSIVE. MONITOR YOUR COSTS*****"+"\033[0m\033[0m")
+
 PINECONE_API_KEY = os.getenv("PINECONE_API_KEY", "")
 assert PINECONE_API_KEY, "PINECONE_API_KEY environment variable is missing from .env"
 
@@ -23,7 +32,7 @@
 assert YOUR_TABLE_NAME, "TABLE_NAME environment variable is missing from .env"
 
 # Project config
-OBJECTIVE = os.getenv("OBJECTIVE", "")
+OBJECTIVE = sys.argv[1] if len(sys.argv) > 1 else os.getenv("OBJECTIVE", "")
 assert OBJECTIVE, "OBJECTIVE environment variable is missing from .env"
 
 YOUR_FIRST_TASK = os.getenv("FIRST_TASK", "")
@@ -58,22 +67,48 @@ def get_ada_embedding(text):
     text = text.replace("\n", " ")
     return openai.Embedding.create(input=[text], model="text-embedding-ada-002")["data"][0]["embedding"]
 
-def task_creation_agent(objective: str, result: Dict, task_description: str, task_list: List[str]):
+def openai_call(prompt: str, use_gpt4: bool = False, temperature: float = 0.5, max_tokens: int = 100):
+    if not use_gpt4:
+        #Call GPT-3 DaVinci model
+        response = openai.Completion.create(
+            engine='text-davinci-003',
+            prompt=prompt,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            top_p=1,
+            frequency_penalty=0,
+            presence_penalty=0
+        )
+        return response.choices[0].text.strip()
+    else:
+        #Call GPT-4 chat model
+        messages=[{"role": "user", "content": prompt}]
+        response = openai.ChatCompletion.create(
+            model="gpt-4",
+            messages = messages,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            n=1,
+            stop=None,
+        )
+        return response.choices[0].message.content.strip()
+
+def task_creation_agent(objective: str, result: Dict, task_description: str, task_list: List[str], gpt_version: str = 'gpt-3'):
     prompt = f"You are an task creation AI that uses the result of an execution agent to create new tasks with the following objective: {objective}, The last completed task has the result: {result}. This result was based on this task description: {task_description}. These are incomplete tasks: {', '.join(task_list)}. Based on the result, create new tasks to be completed by the AI system that do not overlap with incomplete tasks. Return the tasks as an array."
-    response = openai.Completion.create(engine="text-davinci-003",prompt=prompt,temperature=0.5,max_tokens=100,top_p=1,frequency_penalty=0,presence_penalty=0)
-    new_tasks = response.choices[0].text.strip().split('\n')
+    response = openai_call(prompt, USE_GPT4)
+    new_tasks = response.split('\n')
     return [{"task_name": task_name} for task_name in new_tasks]
 
-def prioritization_agent(this_task_id:int):
+def prioritization_agent(this_task_id:int, gpt_version: str = 'gpt-3'):
     global task_list
     task_names = [t["task_name"] for t in task_list]
     next_task_id = int(this_task_id)+1
     prompt = f"""You are an task prioritization AI tasked with cleaning the formatting of and reprioritizing the following tasks: {task_names}. Consider the ultimate objective of your team:{OBJECTIVE}. Do not remove any tasks. Return the result as a numbered list, like:
     #. First task
     #. Second task
     Start the task list with number {next_task_id}."""
-    response = openai.Completion.create(engine="text-davinci-003",prompt=prompt,temperature=0.5,max_tokens=1000,top_p=1,frequency_penalty=0,presence_penalty=0)
-    new_tasks = response.choices[0].text.strip().split('\n')
+    response = openai_call(prompt, USE_GPT4)
+    new_tasks = response.split('\n')
     task_list = deque()
     for task_string in new_tasks:
         task_parts = task_string.strip().split(".", 1)
@@ -82,21 +117,13 @@ def prioritization_agent(this_task_id:int):
             task_name = task_parts[1].strip()
             task_list.append({"task_id": task_id, "task_name": task_name})
 
-def execution_agent(objective:str,task: str) -> str:
+def execution_agent(objective:str,task: str, gpt_version: str = 'gpt-3') -> str:
     #context = context_agent(index="quickstart", query="my_search_query", n=5)
     context=context_agent(index=YOUR_TABLE_NAME, query=objective, n=5)
     #print("\n*******RELEVANT CONTEXT******\n")
     #print(context)
-    response = openai.Completion.create(
-        engine="text-davinci-003",
-        prompt=f"You are an AI who performs one task based on the following objective: {objective}.\nTake into account these previously completed tasks: {context}\nYour task: {task}\nResponse:",
-        temperature=0.7,
-        max_tokens=2000,
-        top_p=1,
-        frequency_penalty=0,
-        presence_penalty=0
-    )
-    return response.choices[0].text.strip()
+    prompt =f"You are an AI who performs one task based on the following objective: {objective}.\nTake into account these previously completed tasks: {context}\nYour task: {task}\nResponse:"
+    return openai_call(prompt, USE_GPT4, 0.7, 2000)
 
 def context_agent(query: str, index: str, n: int):
     query_embedding = get_ada_embedding(query)