outlines

1
"""This example is a simplified translation of BabyAGI.
2

3
It currently does not use the vector store retrieval
4

5
The original repo can be found at https://github.com/yoheinakajima/babyagi
6
"""
7
from collections import deque
8
from typing import Deque, List
9

10
import outlines
11
import outlines.models as models
12

13
model = models.openai("gpt-3.5-turbo")
14
complete = outlines.generate.text(model)
15

16

17
#################
18
# Perform tasks #
19
#################
20

21

22
@outlines.prompt
23
def perform_task_ppt(objective: str, task: str):
24
    """You are an AI who performs one task based on the following objective: {{objective}}.
25

26
    Your task: {{task.task_name}}
27

28
    Response:
29
    """
30

31

32
#####################
33
# Create a new task #
34
#####################
35

36

37
@outlines.prompt
38
def create_tasks_ppt(
39
    objective: str, previous_task: str, result: str, task_list: List[str]
40
):
41
    """You are an task creation AI that uses the result of an execution agent to \
42
    create new tasks with the following objective: {{objective}}.
43

44
    The last completed task has the result: {{result}}.
45

46
    This result was based on this task description: {{previous_task}}. These are \
47
    incomplete tasks: {{task_list | join(task_list)}}.
48

49
    Based on the result, create new tasks to be completed by the AI system that \
50
    do not overlap with incomplete tasks.
51

52
    Return the tasks as an array.
53
    """
54

55

56
def create_tasks_fmt(result: str) -> List[str]:
57
    new_tasks = result.split("\n")
58

59
    task_list = []
60
    for task in new_tasks:
61
        parts = task.strip().split(".", 1)
62
        if len(parts) == 2:
63
            task_list.append(parts[1].strip())
64

65
    return task_list
66

67

68
########################
69
# Prioritize new tasks #
70
########################
71

72

73
@outlines.prompt
74
def prioritize_tasks_ppt(objective: str, task_names: List[str], next_task_id: int):
75
    """You are a task prioritization AI tasked with cleaning the formatting of \
76
    and reprioritizing the following tasks: {{task_names}}.
77

78
    Consider the ultimate objective of your team: {{objective}}.
79

80
    Do not remove any tasks. Return the result as a numbered list, like:
81
    #. First task
82
    #. Second task
83

84
    Start the tasks list with the number {{next_task_id}}.
85
    """
86

87

88
def prioritize_tasks_fmt(result: str):
89
    new_tasks = result.split("\n")
90

91
    task_list: Deque = deque([])
92
    for task in new_tasks:
93
        parts = task.strip().split(".", 1)
94
        if len(parts) == 2:
95
            task_id = int(parts[0].strip())
96
            task_name = parts[1].strip()
97
            task_list.append({"task_id": task_id, "task_name": task_name})
98

99
    return task_list
100

101

102
objective = "Becoming rich while doing nothing."
103
first_task = {
104
    "task_id": 1,
105
    "task_name": "Find a repeatable, low-maintainance, scalable business.",
106
}
107
next_task_id = 1
108
task_list = deque([first_task])
109

110

111
def one_cycle(objective: str, task_list, next_task_id: int):
112
    """One BabyAGI cycle.
113

114
    It consists in executing the highest-priority task, creating some new tasks
115
    given the result, and re-priotizing the tasks.
116

117
    Parameters
118
    ----------
119
    objective
120
        The overall objective of the session.
121
    task_list
122
        The current list of tasks to perform.
123
    task_id_counter
124
        The current task id.
125

126
    """
127

128
    task = task_list.popleft()
129

130
    prompt = perform_task_ppt(objective, task)
131
    result = complete(prompt)
132

133
    prompt = create_tasks_ppt(
134
        objective, first_task["task_name"], result, [first_task["task_name"]]
135
    )
136
    new_tasks = complete(prompt)
137

138
    new_tasks = create_tasks_fmt(new_tasks)
139

140
    for task in new_tasks:
141
        next_task_id += 1
142
        task_list.append({"task_id": next_task_id, "task_name": task})
143

144
    prompt = prioritize_tasks_ppt(
145
        objective, [task["task_name"] for task in task_list], next_task_id
146
    )
147
    prioritized_tasks = complete(prompt)
148

149
    prioritized_tasks = prioritize_tasks_fmt(prioritized_tasks)
150

151
    return task, result, prioritized_tasks, next_task_id
152

153

154
# Let's run it for 5 cycles to see how it works without spending a fortune.
155
for _ in range(5):
156
    print("\033[95m\033[1m" + "\n*****TASK LIST*****\n" + "\033[0m\033[0m")
157
    for t in task_list:
158
        print(" • " + str(t["task_name"]))
159

160
    task, result, task_list, next_task_id = one_cycle(
161
        objective, task_list, next_task_id
162
    )
163

164
    print("\033[92m\033[1m" + "\n*****NEXT TASK*****\n" + "\033[0m\033[0m")
165
    print(task)
166
    print("\033[93m\033[1m" + "\n*****TASK RESULT*****\n" + "\033[0m\033[0m")
167
    print(result)
168