242 lines
8.6 KiB
JavaScript
242 lines
8.6 KiB
JavaScript
import { Configuration, OpenAIApi } from "openai"
|
|
import { ChromaClient, OpenAIEmbeddingFunction } from "chromadb"
|
|
import prompt from "prompt-sync"
|
|
import assert from "assert"
|
|
import * as dotenv from "dotenv"
|
|
dotenv.config()
|
|
|
|
// const client = new ChromaClient("http://localhost:8000")
|
|
|
|
// API Keys
|
|
const OPENAI_API_KEY = process.env.OPENAI_API_KEY || ""
|
|
assert(OPENAI_API_KEY, "OPENAI_API_KEY environment variable is missing from .env")
|
|
|
|
const OPENAI_API_MODEL = process.env.OPENAI_API_MODEL || "gpt-3.5-turbo"
|
|
|
|
// Table config
|
|
const TABLE_NAME = process.env.TABLE_NAME || ""
|
|
assert(TABLE_NAME, "TABLE_NAME environment variable is missing from .env")
|
|
|
|
// Run config
|
|
const BABY_NAME = process.env.BABY_NAME || "BabyAGI"
|
|
|
|
// Goal config
|
|
const p = prompt()
|
|
const OBJECTIVE = p("What is BabyAGI's objective? ")
|
|
const INITIAL_TASK = p("What is the initial task to complete the objective? ")
|
|
assert(OBJECTIVE, "No objective provided.")
|
|
assert (INITIAL_TASK, "No initial task provided.")
|
|
|
|
console.log('\x1b[95m\x1b[1m\n*****CONFIGURATION*****\n\x1b[0m\x1b[0m')
|
|
console.log(`Name: ${BABY_NAME}`)
|
|
console.log(`LLM: ${OPENAI_API_MODEL}`)
|
|
|
|
if (OPENAI_API_MODEL.toLowerCase().includes("gpt-4")){
|
|
console.log("\x1b[91m\x1b[1m\n*****USING GPT-4. POTENTIALLY EXPENSIVE. MONITOR YOUR COSTS*****\x1b[0m\x1b[0m")
|
|
}
|
|
|
|
console.log("\x1b[94m\x1b[1m" + "\n*****OBJECTIVE*****\n" + "\x1b[0m\x1b[0m")
|
|
console.log(`${OBJECTIVE}`)
|
|
|
|
console.log(`\x1b[93m\x1b[1m \nInitial task: \x1b[0m\x1b[0m ${INITIAL_TASK}`)
|
|
|
|
// Define OpenAI embedding function using Chroma
|
|
const embeddingFunction = new OpenAIEmbeddingFunction(OPENAI_API_KEY)
|
|
|
|
// Configure OpenAI
|
|
const configuration = new Configuration({
|
|
apiKey: OPENAI_API_KEY,
|
|
});
|
|
const openai = new OpenAIApi(configuration);
|
|
|
|
//Task List
|
|
var taskList = []
|
|
|
|
// Connect to chromadb and create/get collection
|
|
const chromaConnect = async ()=>{
|
|
const chroma = new ChromaClient("http://localhost:8000")
|
|
const metric = "cosine"
|
|
const collections = await chroma.listCollections()
|
|
const collectionNames = collections.map((c)=>c.name)
|
|
if(collectionNames.includes(TABLE_NAME)){
|
|
const collection = await chroma.getCollection(TABLE_NAME, embeddingFunction)
|
|
return collection
|
|
}
|
|
else{
|
|
const collection = await chroma.createCollection(
|
|
TABLE_NAME,
|
|
{
|
|
"hnsw:space": metric
|
|
},
|
|
embeddingFunction
|
|
)
|
|
return collection
|
|
}
|
|
}
|
|
|
|
const add_task = (task)=>{ taskList.push(task) }
|
|
|
|
const clear_tasks = ()=>{ taskList = [] }
|
|
|
|
const get_ada_embedding = async (text)=>{
|
|
text = text.replace("\n", " ")
|
|
const embedding = await embeddingFunction.generate(text)
|
|
return embedding
|
|
}
|
|
|
|
const openai_completion = async (prompt, temperature=0.5, maxTokens=100)=>{
|
|
if(OPENAI_API_MODEL.startsWith("gpt-")){
|
|
const messages = [{"role": "system", "content": prompt}]
|
|
const response = await openai.createChatCompletion({
|
|
model: OPENAI_API_MODEL,
|
|
messages: messages,
|
|
max_tokens: maxTokens,
|
|
temperature: temperature,
|
|
n: 1,
|
|
stop: null
|
|
})
|
|
return response.data.choices[0].message.content.trim()
|
|
}
|
|
else {
|
|
const response = await openai.createCompletion({
|
|
model: OPENAI_API_MODEL,
|
|
prompt: prompt,
|
|
max_tokens: maxTokens,
|
|
temperature: temperature,
|
|
top_p: 1,
|
|
frequency_penalty: 0,
|
|
presence_penalty: 0
|
|
})
|
|
return response.data.choices[0].text.trim()
|
|
}
|
|
}
|
|
|
|
const task_creation_agent = async (objective, result, task_description, taskList)=>{
|
|
const prompt = `
|
|
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: ${taskList.map(task=>`${task.taskId}: ${task.taskName}`).join(', ')}.
|
|
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.`
|
|
const response = await openai_completion(prompt)
|
|
const newTasks = response.trim().includes("\n") ? response.trim().split("\n") : [response.trim()];
|
|
return newTasks.map(taskName => ({ taskName: taskName }));
|
|
}
|
|
|
|
|
|
|
|
const prioritization_agent = async (taskId)=>{
|
|
const taskNames = taskList.map((task)=>task.taskName)
|
|
const nextTaskId = taskId+1
|
|
const prompt = `
|
|
You are an task prioritization AI tasked with cleaning the formatting of and reprioritizing the following tasks: ${taskNames}.
|
|
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 ${nextTaskId}.`
|
|
const response = await openai_completion(prompt)
|
|
const newTasks = response.trim().includes("\n") ? response.trim().split("\n") : [response.trim()];
|
|
clear_tasks()
|
|
newTasks.forEach((newTask)=>{
|
|
const newTaskParts = newTask.trim().split(/\.(?=\s)/)
|
|
if (newTaskParts.length == 2){
|
|
const newTaskId = newTaskParts[0].trim()
|
|
const newTaskName = newTaskParts[1].trim()
|
|
add_task({
|
|
taskId: newTaskId,
|
|
taskName: newTaskName
|
|
})
|
|
}
|
|
})
|
|
}
|
|
|
|
const execution_agent = async (objective, task, chromaCollection)=>{
|
|
const context = context_agent(objective, 5, chromaCollection)
|
|
const prompt = `
|
|
You are an AI who performs one task based on the following objective: ${objective}.\n
|
|
Take into account these previously completed tasks: ${context}.\n
|
|
Your task: ${task}\nResponse:`
|
|
const response = await openai_completion(prompt, undefined, 2000)
|
|
return response
|
|
}
|
|
|
|
const context_agent = async (query, topResultsNum, chromaCollection)=>{
|
|
const count = await chromaCollection.count()
|
|
if (count == 0){
|
|
return []
|
|
}
|
|
const results = await chromaCollection.query(
|
|
undefined,
|
|
Math.min(topResultsNum, count),
|
|
undefined,
|
|
query,
|
|
)
|
|
return results.metadatas[0].map(item=>item.task)
|
|
}
|
|
|
|
function sleep(ms) {
|
|
return new Promise(resolve => setTimeout(resolve, ms))
|
|
}
|
|
|
|
(async()=>{
|
|
const initialTask = {
|
|
taskId: 1,
|
|
taskName: INITIAL_TASK
|
|
}
|
|
add_task(initialTask)
|
|
const chromaCollection = await chromaConnect()
|
|
var taskIdCounter = 1
|
|
while (true){
|
|
if(taskList.length>0){
|
|
console.log("\x1b[95m\x1b[1m"+"\n*****TASK LIST*****\n"+"\x1b[0m\x1b[0m")
|
|
taskList.forEach(t => {
|
|
console.log(" • " + t.taskName)
|
|
})
|
|
|
|
// Step 1: Pull the first task
|
|
const task = taskList.shift()
|
|
console.log("\x1b[92m\x1b[1m"+"\n*****NEXT TASK*****\n"+"\x1b[0m\x1b[0m")
|
|
console.log(task.taskId + ": " + task.taskName)
|
|
|
|
// Send to execution function to complete the task based on the context
|
|
const result = await execution_agent(OBJECTIVE, task.taskName, chromaCollection)
|
|
const currTaskId = task.taskId
|
|
console.log("\x1b[93m\x1b[1m"+"\nTASK RESULT\n"+"\x1b[0m\x1b[0m")
|
|
console.log(result)
|
|
|
|
// Step 2: Enrich result and store in Chroma
|
|
const enrichedResult = { data : result} // this is where you should enrich the result if needed
|
|
const resultId = `result_${task.taskId}`
|
|
const vector = enrichedResult.data // extract the actual result from the dictionary
|
|
const collectionLength = (await chromaCollection.get([resultId])).ids?.length
|
|
if(collectionLength>0){
|
|
await chromaCollection.update(
|
|
resultId,
|
|
undefined,
|
|
{task: task.taskName, result: result},
|
|
vector
|
|
)
|
|
}
|
|
else{
|
|
await chromaCollection.add(
|
|
resultId,
|
|
undefined,
|
|
{task: task.taskName, result},
|
|
vector
|
|
)
|
|
}
|
|
|
|
// Step 3: Create new tasks and reprioritize task list
|
|
const newTasks = await task_creation_agent(OBJECTIVE, enrichedResult, task.taskName, taskList.map(task=>task.taskName))
|
|
newTasks.forEach((task)=>{
|
|
taskIdCounter += 1
|
|
task.taskId = taskIdCounter
|
|
add_task(task)
|
|
})
|
|
await prioritization_agent(currTaskId)
|
|
await sleep(3000)
|
|
}
|
|
}
|
|
})()
|
|
|