kenken999's picture
updatest
14dc68f
raw
history blame
8.83 kB
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)
}
}
})()