claude's babyagi implementation

claudes-babyagi.py

import os
from collections import deque
from typing import Dict, List
import time
import openai

# Configuration
OBJECTIVE = "Solve world hunger"
YOUR_API_KEY = "your-openai-api-key-here"
INITIAL_TASK = "Develop a comprehensive plan to solve world hunger"

# Set up OpenAI API
openai.api_key = YOUR_API_KEY

class Task:
    def __init__(self, task_id: int, description: str):
        self.task_id = task_id
        self.description = description

def get_ada_embedding(text):
    text = text.replace("\n", " ")
    return openai.Embedding.create(input=[text], model="text-embedding-ada-002")["data"][0]["embedding"]

def openai_call(prompt: str, model: str = "gpt-3.5-turbo", temperature: float = 0.5, max_tokens: int = 100):
    response = openai.ChatCompletion.create(
        model=model,
        messages=[{"role": "system", "content": "You are an AI assistant."},
                  {"role": "user", "content": prompt}],
        temperature=temperature,
        max_tokens=max_tokens
    )
    return response.choices[0].message.content.strip()

def task_creation_agent(objective: str, result: Dict, task_description: str, task_list: List[str]):
    prompt = f"Objective: {objective}\n"
    prompt += f"Result: {result}\n"
    prompt += f"Task Description: {task_description}\n"
    prompt += f"Current Task List: {', '.join(task_list)}\n"
    prompt += "Based on the result, create new tasks to be completed in order to meet the objective. Tasks should be simple and actionable."

    response = openai_call(prompt)
    new_tasks = response.split('\n')
    return [task.strip() for task in new_tasks if task.strip()]

def prioritization_agent(task_id: int, task_list: List[Dict]):
    task_names = [t["task_name"] for t in task_list]
    prompt = f"Task ID: {task_id}\n"
    prompt += f"Task List: {task_names}\n"
    prompt += "Prioritize these tasks based on their importance for achieving the objective. Return the tasks as a comma-separated list, with the highest priority task first."

    response = openai_call(prompt)
    prioritized_tasks = response.split(',')
    return [task.strip() for task in prioritized_tasks if task.strip()]

def execution_agent(objective: str, task: str) -> str:
    prompt = f"Objective: {objective}\n"
    prompt += f"Task: {task}\n"
    prompt += "Execute the task and provide the result."

    return openai_call(prompt, max_tokens=2000)

def main():
    task_list = deque([])
    task_id_counter = 0
    task_list.append(Task(task_id_counter, INITIAL_TASK))

    while True:
        if not task_list:
            print("All tasks completed.")
            break

        current_task = task_list.popleft()
        print(f"\nExecuting Task {current_task.task_id}: {current_task.description}")

        result = execution_agent(OBJECTIVE, current_task.description)
        print(f"\nTask Result: {result}")

        new_tasks = task_creation_agent(OBJECTIVE, result, current_task.description, [t.description for t in task_list])
        
        for new_task in new_tasks:
            task_id_counter += 1
            task_list.append(Task(task_id_counter, new_task))

        if task_list:
            prioritized_tasks = prioritization_agent(current_task.task_id, [{"task_id": t.task_id, "task_name": t.description} for t in task_list])
            task_list = deque(sorted(task_list, key=lambda x: prioritized_tasks.index(x.description) if x.description in prioritized_tasks else len(prioritized_tasks)))

        time.sleep(1)  # Pause to avoid hitting API rate limits

if __name__ == "__main__":
    main()