aljaz/Auto-GPT
1
0
Fork 0
mirror of https://github.com/aljazceru/Auto-GPT.git synced 2025-12-24 01:14:22 +01:00
Code Issues Packages Projects Releases Wiki Activity

Make Auto-GPT aware of its running cost (#762)

Browse Source 
* Implemented running cost counter for chat completions

This data is known to the AI as additional system context, and is printed out to the user

* Added comments to api_manager.py

* Added user-defined API budget.

The user is now prompted if they want to give the AI a budget for API calls. If they enter nothing, there is no monetary limit, but if they define a budget then the AI will be told to shut down gracefully once it has come within 1 cent of its limit, and to shut down immediately once it has exceeded its limit. If a budget is defined, Auto-GPT is always aware of how much it was given and how much remains to be spent.

* Chat completion calls are now done through api_manager. Total running cost is printed.

* Implemented api budget setting and tracking

User can now configure a maximum api budget, and the AI is aware of that and its remaining budget. The AI is instructed to shut down when exceeding the budget.

* Update autogpt/api_manager.py

Change "per token" to "per 1000 tokens" in a comment on the api cost

Co-authored-by: Rob Luke <code@robertluke.net>

* Fixed lint errors

* Include embedding costs

* Add embedding completion cost

* lint

* Added 'requires_api_key' decorator to test_commands.py, switched to a valid chat completions model

* Refactor API manager, add debug mode, and add tests

- Extract model costs to  to avoid duplication
- Add debug mode parameter to ApiManager class
- Move debug mode configuration to
- Log AI response and budget messages in debug mode
- Implement 'test_api_manager.py'

* Fixed test_setup failing. An extra user input is needed for api budget

* Linting

---------

Co-authored-by: Rob Luke <code@robertluke.net>
Co-authored-by: Nicholas Tindle <nick@ntindle.com>
This commit is contained in:
Vwing
2023-04-23 14:04:31 -07:00
committed by GitHub
parent bf895eb656
commit d6ef9d1b5d
11 changed files with 401 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

158
autogpt/api_manager.py Normal file
View File

@@ -0,0 +1,158 @@
from typing import List
import openai
from autogpt.config import Config
from autogpt.logs import logger
from autogpt.modelsinfo import COSTS
cfg = Config()
openai.api_key = cfg.openai_api_key
print_total_cost = cfg.debug_mode
class ApiManager:
def __init__(self, debug=False):
self.total_prompt_tokens = 0
self.total_completion_tokens = 0
self.total_cost = 0
self.total_budget = 0
self.debug = debug
def reset(self):
self.total_prompt_tokens = 0
self.total_completion_tokens = 0
self.total_cost = 0
self.total_budget = 0.0
def create_chat_completion(
self,
messages: list, # type: ignore
model: str | None = None,
temperature: float = cfg.temperature,
max_tokens: int | None = None,
deployment_id=None,
) -> str:
"""
Create a chat completion and update the cost.
Args:
messages (list): The list of messages to send to the API.
model (str): The model to use for the API call.
temperature (float): The temperature to use for the API call.
max_tokens (int): The maximum number of tokens for the API call.
Returns:
str: The AI's response.
"""
if deployment_id is not None:
response = openai.ChatCompletion.create(
deployment_id=deployment_id,
model=model,
messages=messages,
temperature=temperature,
max_tokens=max_tokens,
)
else:
response = openai.ChatCompletion.create(
model=model,
messages=messages,
temperature=temperature,
max_tokens=max_tokens,
)
if self.debug:
logger.debug(f"Response: {response}")
prompt_tokens = response.usage.prompt_tokens
completion_tokens = response.usage.completion_tokens
self.update_cost(prompt_tokens, completion_tokens, model)
return response
def embedding_create(
self,
text_list: List[str],
model: str = "text-embedding-ada-002",
) -> List[float]:
"""
Create an embedding for the given input text using the specified model.
Args:
text_list (List[str]): Input text for which the embedding is to be created.
model (str, optional): The model to use for generating the embedding.
Returns:
List[float]: The generated embedding as a list of float values.
"""
if cfg.use_azure:
response = openai.Embedding.create(
input=text_list,
engine=cfg.get_azure_deployment_id_for_model(model),
)
else:
response = openai.Embedding.create(input=text_list, model=model)
self.update_cost(response.usage.prompt_tokens, 0, model)
return response["data"][0]["embedding"]
def update_cost(self, prompt_tokens, completion_tokens, model):
"""
Update the total cost, prompt tokens, and completion tokens.
Args:
prompt_tokens (int): The number of tokens used in the prompt.
completion_tokens (int): The number of tokens used in the completion.
model (str): The model used for the API call.
"""
self.total_prompt_tokens += prompt_tokens
self.total_completion_tokens += completion_tokens
self.total_cost += (
prompt_tokens * COSTS[model]["prompt"]
+ completion_tokens * COSTS[model]["completion"]
) / 1000
if print_total_cost:
print(f"Total running cost: ${self.total_cost:.3f}")
def set_total_budget(self, total_budget):
"""
Sets the total user-defined budget for API calls.
Args:
prompt_tokens (int): The number of tokens used in the prompt.
"""
self.total_budget = total_budget
def get_total_prompt_tokens(self):
"""
Get the total number of prompt tokens.
Returns:
int: The total number of prompt tokens.
"""
return self.total_prompt_tokens
def get_total_completion_tokens(self):
"""
Get the total number of completion tokens.
Returns:
int: The total number of completion tokens.
"""
return self.total_completion_tokens
def get_total_cost(self):
"""
Get the total cost of API calls.
Returns:
float: The total cost of API calls.
"""
return self.total_cost
def get_total_budget(self):
"""
Get the total user-defined budget for API calls.
Returns:
float: The total budget for API calls.
"""
return self.total_budget
api_manager = ApiManager(cfg.debug_mode)

23
autogpt/chat.py
View File

@@ -3,6 +3,7 @@ import time
from openai.error import RateLimitError
from autogpt import token_counter
from autogpt.api_manager import api_manager
from autogpt.config import Config
from autogpt.llm_utils import create_chat_completion
from autogpt.logs import logger
@@ -133,6 +134,28 @@ def chat_with_ai(
# Move to the next most recent message in the full message history
next_message_to_add_index -= 1
# inform the AI about its remaining budget (if it has one)
if api_manager.get_total_budget() > 0.0:
remaining_budget = (
api_manager.get_total_budget() - api_manager.get_total_cost()
)
if remaining_budget < 0:
remaining_budget = 0
system_message = (
f"Your remaining API budget is ${remaining_budget:.3f}"
+ (
" BUDGET EXCEEDED! SHUT DOWN!\n\n"
if remaining_budget == 0
else " Budget very nearly exceeded! Shut down gracefully!\n\n"
if remaining_budget < 0.005
else " Budget nearly exceeded. Finish up.\n\n"
if remaining_budget < 0.01
else "\n\n"
)
)
logger.debug(system_message)
current_context.append(create_chat_message("system", system_message))
# Append user input, the length of this is accounted for above
current_context.extend([create_chat_message("user", user_input)])

19
autogpt/config/ai_config.py
View File

@@ -26,10 +26,15 @@ class AIConfig:
ai_name (str): The name of the AI.
ai_role (str): The description of the AI's role.
ai_goals (list): The list of objectives the AI is supposed to complete.
api_budget (float): The maximum dollar value for API calls (0.0 means infinite)
"""
def __init__(
self, ai_name: str = "", ai_role: str = "", ai_goals: list | None = None
self,
ai_name: str = "",
ai_role: str = "",
ai_goals: list | None = None,
api_budget: float = 0.0,
) -> None:
"""
Initialize a class instance
@@ -38,6 +43,7 @@ class AIConfig:
ai_name (str): The name of the AI.
ai_role (str): The description of the AI's role.
ai_goals (list): The list of objectives the AI is supposed to complete.
api_budget (float): The maximum dollar value for API calls (0.0 means infinite)
Returns:
None
"""
@@ -46,13 +52,14 @@ class AIConfig:
self.ai_name = ai_name
self.ai_role = ai_role
self.ai_goals = ai_goals
self.api_budget = api_budget
self.prompt_generator = None
self.command_registry = None
@staticmethod
def load(config_file: str = SAVE_FILE) -> "AIConfig":
"""
Returns class object with parameters (ai_name, ai_role, ai_goals) loaded from
Returns class object with parameters (ai_name, ai_role, ai_goals, api_budget) loaded from
yaml file if yaml file exists,
else returns class with no parameters.
@@ -73,8 +80,9 @@ class AIConfig:
ai_name = config_params.get("ai_name", "")
ai_role = config_params.get("ai_role", "")
ai_goals = config_params.get("ai_goals", [])
api_budget = config_params.get("api_budget", 0.0)
# type: Type[AIConfig]
return AIConfig(ai_name, ai_role, ai_goals)
return AIConfig(ai_name, ai_role, ai_goals, api_budget)
def save(self, config_file: str = SAVE_FILE) -> None:
"""
@@ -92,6 +100,7 @@ class AIConfig:
"ai_name": self.ai_name,
"ai_role": self.ai_role,
"ai_goals": self.ai_goals,
"api_budget": self.api_budget,
}
with open(config_file, "w", encoding="utf-8") as file:
yaml.dump(config, file, allow_unicode=True)
@@ -107,7 +116,7 @@ class AIConfig:
Returns:
full_prompt (str): A string containing the initial prompt for the user
including the ai_name, ai_role and ai_goals.
including the ai_name, ai_role, ai_goals, and api_budget.
"""
prompt_start = (
@@ -147,6 +156,8 @@ class AIConfig:
full_prompt = f"You are {prompt_generator.name}, {prompt_generator.role}\n{prompt_start}\n\nGOALS:\n\n"
for i, goal in enumerate(self.ai_goals):
full_prompt += f"{i+1}. {goal}\n"
if self.api_budget > 0.0:
full_prompt += f"\nIt takes money to let you run. Your API budget is ${self.api_budget:.3f}"
self.prompt_generator = prompt_generator
full_prompt += f"\n\n{prompt_generator.generate_prompt_string()}"
return full_prompt

19
autogpt/llm_utils.py
View File

@@ -7,6 +7,7 @@ import openai
from colorama import Fore, Style
from openai.error import APIError, RateLimitError
from autogpt.api_manager import api_manager
from autogpt.config import Config
from autogpt.logs import logger
from autogpt.types.openai import Message
@@ -96,7 +97,7 @@ def create_chat_completion(
backoff = 2 ** (attempt + 2)
try:
if CFG.use_azure:
response = openai.ChatCompletion.create(
response = api_manager.create_chat_completion(
deployment_id=CFG.get_azure_deployment_id_for_model(model),
model=model,
messages=messages,
@@ -104,7 +105,7 @@ def create_chat_completion(
max_tokens=max_tokens,
)
else:
response = openai.ChatCompletion.create(
response = api_manager.create_chat_completion(
model=model,
messages=messages,
temperature=temperature,
@@ -159,17 +160,9 @@ def create_embedding_with_ada(text) -> list:
for attempt in range(num_retries):
backoff = 2 ** (attempt + 2)
try:
if CFG.use_azure:
return openai.Embedding.create(
input=[text],
engine=CFG.get_azure_deployment_id_for_model(
"text-embedding-ada-002"
),
)["data"][0]["embedding"]
else:
return openai.Embedding.create(
input=[text], model="text-embedding-ada-002"
)["data"][0]["embedding"]
return api_manager.embedding_create(
text_list=[text], model="text-embedding-ada-002"
)
except RateLimitError:
pass
except APIError as e:

13
autogpt/memory/base.py
View File

@@ -3,6 +3,7 @@ import abc
import openai
from autogpt.api_manager import api_manager
from autogpt.config import AbstractSingleton, Config
cfg = Config()
@@ -10,15 +11,9 @@ cfg = Config()
def get_ada_embedding(text):
text = text.replace("\n", " ")
if cfg.use_azure:
return openai.Embedding.create(
input=[text],
engine=cfg.get_azure_deployment_id_for_model("text-embedding-ada-002"),
)["data"][0]["embedding"]
else:
return openai.Embedding.create(input=[text], model="text-embedding-ada-002")[
"data"
][0]["embedding"]
return api_manager.embedding_create(
text_list=[text], model="text-embedding-ada-002"
)
class MemoryProviderSingleton(AbstractSingleton):

7
autogpt/modelsinfo.py Normal file
View File

@@ -0,0 +1,7 @@
COSTS = {
"gpt-3.5-turbo": {"prompt": 0.002, "completion": 0.002},
"gpt-3.5-turbo-0301": {"prompt": 0.002, "completion": 0.002},
"gpt-4-0314": {"prompt": 0.03, "completion": 0.06},
"gpt-4": {"prompt": 0.03, "completion": 0.06},
"text-embedding-ada-002": {"prompt": 0.0004, "completion": 0.0},
}

10
autogpt/prompts/prompt.py
View File

@@ -1,5 +1,6 @@
from colorama import Fore
from autogpt.api_manager import api_manager
from autogpt.config.ai_config import AIConfig
from autogpt.config.config import Config
from autogpt.logs import logger
@@ -86,6 +87,11 @@ def construct_main_ai_config() -> AIConfig:
logger.typewriter_log("Name :", Fore.GREEN, config.ai_name)
logger.typewriter_log("Role :", Fore.GREEN, config.ai_role)
logger.typewriter_log("Goals:", Fore.GREEN, f"{config.ai_goals}")
logger.typewriter_log(
"API Budget:",
Fore.GREEN,
"infinite" if config.api_budget <= 0 else f"${config.api_budget}",
)
elif config.ai_name:
logger.typewriter_log(
"Welcome back! ",
@@ -98,6 +104,7 @@ def construct_main_ai_config() -> AIConfig:
Name: {config.ai_name}
Role: {config.ai_role}
Goals: {config.ai_goals}
API Budget: {"infinite" if config.api_budget <= 0 else f"${config.api_budget}"}
Continue (y/n): """
)
if should_continue.lower() == "n":
@@ -107,6 +114,9 @@ Continue (y/n): """
config = prompt_user()
config.save(CFG.ai_settings_file)
# set the total api budget
api_manager.set_total_budget(config.api_budget)
# Agent Created, print message
logger.typewriter_log(
config.ai_name,

26
autogpt/setup.py
View File

@@ -133,7 +133,28 @@ def generate_aiconfig_manual() -> AIConfig:
"Develop and manage multiple businesses autonomously",
]
return AIConfig(ai_name, ai_role, ai_goals)
# Get API Budget from User
logger.typewriter_log(
"Enter your budget for API calls: ",
Fore.GREEN,
"For example: $1.50",
)
print("Enter nothing to let the AI run without monetary limit", flush=True)
api_budget_input = utils.clean_input(
f"{Fore.LIGHTBLUE_EX}Budget{Style.RESET_ALL}: $"
)
if api_budget_input == "":
api_budget = 0.0
else:
try:
api_budget = float(api_budget_input.replace("$", ""))
except ValueError:
logger.typewriter_log(
"Invalid budget input. Setting budget to unlimited.", Fore.RED
)
api_budget = 0.0
return AIConfig(ai_name, ai_role, ai_goals, api_budget)
def generate_aiconfig_automatic(user_prompt) -> AIConfig:
@@ -192,5 +213,6 @@ Goals:
.strip()
)
ai_goals = re.findall(r"(?<=\n)-\s*(.*)", output)
api_budget = 0.0 # TODO: parse api budget using a regular expression
return AIConfig(ai_name, ai_role, ai_goals)
return AIConfig(ai_name, ai_role, ai_goals, api_budget)

148
tests/test_api_manager.py Normal file
View File

@@ -0,0 +1,148 @@
from unittest.mock import MagicMock, patch
import pytest
from autogpt.api_manager import COSTS, ApiManager
api_manager = ApiManager()
@pytest.fixture(autouse=True)
def reset_api_manager():
api_manager.reset()
yield
@pytest.fixture(autouse=True)
def mock_costs():
with patch.dict(
COSTS,
{
"gpt-3.5-turbo": {"prompt": 0.002, "completion": 0.002},
"text-embedding-ada-002": {"prompt": 0.0004, "completion": 0},
},
clear=True,
):
yield
class TestApiManager:
@staticmethod
def test_create_chat_completion_debug_mode(caplog):
"""Test if debug mode logs response."""
api_manager_debug = ApiManager(debug=True)
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Who won the world series in 2020?"},
]
model = "gpt-3.5-turbo"
with patch("openai.ChatCompletion.create") as mock_create:
mock_response = MagicMock()
mock_response.usage.prompt_tokens = 10
mock_response.usage.completion_tokens = 20
mock_create.return_value = mock_response
api_manager_debug.create_chat_completion(messages, model=model)
assert "Response" in caplog.text
@staticmethod
def test_create_chat_completion_empty_messages():
"""Test if empty messages result in zero tokens and cost."""
messages = []
model = "gpt-3.5-turbo"
with patch("openai.ChatCompletion.create") as mock_create:
mock_response = MagicMock()
mock_response.usage.prompt_tokens = 0
mock_response.usage.completion_tokens = 0
mock_create.return_value = mock_response
api_manager.create_chat_completion(messages, model=model)
assert api_manager.get_total_prompt_tokens() == 0
assert api_manager.get_total_completion_tokens() == 0
assert api_manager.get_total_cost() == 0
@staticmethod
def test_create_chat_completion_valid_inputs():
"""Test if valid inputs result in correct tokens and cost."""
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Who won the world series in 2020?"},
]
model = "gpt-3.5-turbo"
with patch("openai.ChatCompletion.create") as mock_create:
mock_response = MagicMock()
mock_response.usage.prompt_tokens = 10
mock_response.usage.completion_tokens = 20
mock_create.return_value = mock_response
api_manager.create_chat_completion(messages, model=model)
assert api_manager.get_total_prompt_tokens() == 10
assert api_manager.get_total_completion_tokens() == 20
assert api_manager.get_total_cost() == (10 * 0.002 + 20 * 0.002) / 1000
@staticmethod
def test_embedding_create_invalid_model():
"""Test if an invalid model for embedding raises a KeyError."""
text_list = ["Hello, how are you?"]
model = "invalid-model"
with patch("openai.Embedding.create") as mock_create:
mock_response = MagicMock()
mock_response.usage.prompt_tokens = 5
mock_create.side_effect = KeyError("Invalid model")
with pytest.raises(KeyError):
api_manager.embedding_create(text_list, model=model)
@staticmethod
def test_embedding_create_valid_inputs():
"""Test if valid inputs for embedding result in correct tokens and cost."""
text_list = ["Hello, how are you?"]
model = "text-embedding-ada-002"
with patch("openai.Embedding.create") as mock_create:
mock_response = MagicMock()
mock_response.usage.prompt_tokens = 5
mock_response["data"] = [{"embedding": [0.1, 0.2, 0.3]}]
mock_create.return_value = mock_response
api_manager.embedding_create(text_list, model=model)
assert api_manager.get_total_prompt_tokens() == 5
assert api_manager.get_total_completion_tokens() == 0
assert api_manager.get_total_cost() == (5 * 0.0004) / 1000
def test_getter_methods(self):
"""Test the getter methods for total tokens, cost, and budget."""
api_manager.update_cost(60, 120, "gpt-3.5-turbo")
api_manager.set_total_budget(10.0)
assert api_manager.get_total_prompt_tokens() == 60
assert api_manager.get_total_completion_tokens() == 120
assert api_manager.get_total_cost() == (60 * 0.002 + 120 * 0.002) / 1000
assert api_manager.get_total_budget() == 10.0
@staticmethod
def test_set_total_budget():
"""Test if setting the total budget works correctly."""
total_budget = 10.0
api_manager.set_total_budget(total_budget)
assert api_manager.get_total_budget() == total_budget
@staticmethod
def test_update_cost():
"""Test if updating the cost works correctly."""
prompt_tokens = 50
completion_tokens = 100
model = "gpt-3.5-turbo"
api_manager.update_cost(prompt_tokens, completion_tokens, model)
assert api_manager.get_total_prompt_tokens() == 50
assert api_manager.get_total_completion_tokens() == 100
assert api_manager.get_total_cost() == (50 * 0.002 + 100 * 0.002) / 1000

6
tests/unit/test_commands.py
View File

@@ -5,18 +5,20 @@ import pytest
import autogpt.agent.agent_manager as agent_manager
from autogpt.app import execute_command, list_agents, start_agent
from tests.utils import requires_api_key
@pytest.mark.integration_test
@requires_api_key("OPENAI_API_KEY")
def test_make_agent() -> None:
"""Test the make_agent command"""
with patch("openai.ChatCompletion.create") as mock:
obj = MagicMock()
obj.response.choices[0].messages[0].content = "Test message"
mock.return_value = obj
start_agent("Test Agent", "chat", "Hello, how are you?", "gpt2")
start_agent("Test Agent", "chat", "Hello, how are you?", "gpt-3.5-turbo")
agents = list_agents()
assert "List of agents:\n0: chat" == agents
start_agent("Test Agent 2", "write", "Hello, how are you?", "gpt2")
start_agent("Test Agent 2", "write", "Hello, how are you?", "gpt-3.5-turbo")
agents = list_agents()
assert "List of agents:\n0: chat\n1: write" == agents

2
tests/unit/test_setup.py
View File

@@ -44,6 +44,7 @@ class TestAutoGPT(unittest.TestCase):
"Purchase ingredients",
"Bake a cake",
"",
"",
]
with patch("builtins.input", side_effect=user_inputs):
ai_config = prompt_user()
@@ -62,6 +63,7 @@ class TestAutoGPT(unittest.TestCase):
"Purchase ingredients",
"Bake a cake",
"",
"",
]
with patch("builtins.input", side_effect=user_inputs):
ai_config = prompt_user()