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

* 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 <[email protected]>

* 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 <[email protected]>
Co-authored-by: Nicholas Tindle <[email protected]>

autogpt/api_manager.py

@@ -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)

autogpt/chat.py

@@ -3,6 +3,7 @@
 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)])
 

autogpt/config/ai_config.py

@@ -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 @@ def __init__(
             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 @@ def __init__(
         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 @@ def load(config_file: str = SAVE_FILE) -> "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 @@ def save(self, config_file: str = SAVE_FILE) -> None:
             "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 @@ def construct_full_prompt(
 
         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 @@ def construct_full_prompt(
         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

autogpt/llm_utils.py

@@ -7,6 +7,7 @@
 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,15 +97,15 @@ 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,
                     temperature=temperature,
                     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:

autogpt/memory/base.py

@@ -3,22 +3,17 @@
 
 import openai
 
+from autogpt.api_manager import api_manager
 from autogpt.config import AbstractSingleton, Config
 
 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):

autogpt/modelsinfo.py

@@ -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},
+}

autogpt/prompts/prompt.py

@@ -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 @@ def construct_main_ai_config() -> AIConfig:
         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,