diff --git a/autogpt/api_manager.py b/autogpt/api_manager.py index ace64fbd..52b94ff7 100644 --- a/autogpt/api_manager.py +++ b/autogpt/api_manager.py @@ -1,22 +1,19 @@ -from typing import List +from __future__ import annotations import openai from autogpt.config import Config from autogpt.logs import logger from autogpt.modelsinfo import COSTS - -cfg = Config() -print_total_cost = cfg.debug_mode +from autogpt.singleton import Singleton -class ApiManager: - def __init__(self, debug=False): +class ApiManager(metaclass=Singleton): + def __init__(self): 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 @@ -28,7 +25,7 @@ class ApiManager: self, messages: list, # type: ignore model: str | None = None, - temperature: float = cfg.temperature, + temperature: float = None, max_tokens: int | None = None, deployment_id=None, ) -> str: @@ -42,6 +39,9 @@ class ApiManager: Returns: str: The AI's response. """ + cfg = Config() + if temperature is None: + temperature = cfg.temperature if deployment_id is not None: response = openai.ChatCompletion.create( deployment_id=deployment_id, @@ -59,8 +59,7 @@ class ApiManager: max_tokens=max_tokens, api_key=cfg.openai_api_key, ) - if self.debug: - logger.debug(f"Response: {response}") + 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) @@ -81,8 +80,7 @@ class ApiManager: prompt_tokens * COSTS[model]["prompt"] + completion_tokens * COSTS[model]["completion"] ) / 1000 - if print_total_cost: - print(f"Total running cost: ${self.total_cost:.3f}") + logger.debug(f"Total running cost: ${self.total_cost:.3f}") def set_total_budget(self, total_budget): """ @@ -128,6 +126,3 @@ class ApiManager: float: The total budget for API calls. """ return self.total_budget - - -api_manager = ApiManager(cfg.debug_mode) diff --git a/autogpt/chat.py b/autogpt/chat.py index 21eab6a0..4b906a00 100644 --- a/autogpt/chat.py +++ b/autogpt/chat.py @@ -3,7 +3,7 @@ import time from openai.error import RateLimitError from autogpt import token_counter -from autogpt.api_manager import api_manager +from autogpt.api_manager import ApiManager from autogpt.config import Config from autogpt.llm_utils import create_chat_completion from autogpt.logs import logger @@ -134,6 +134,7 @@ def chat_with_ai( # Move to the next most recent message in the full message history next_message_to_add_index -= 1 + api_manager = ApiManager() # inform the AI about its remaining budget (if it has one) if api_manager.get_total_budget() > 0.0: remaining_budget = ( diff --git a/autogpt/llm_utils.py b/autogpt/llm_utils.py index 212928ca..a98b12a3 100644 --- a/autogpt/llm_utils.py +++ b/autogpt/llm_utils.py @@ -8,7 +8,7 @@ import openai from colorama import Fore, Style from openai.error import APIError, RateLimitError, Timeout -from autogpt.api_manager import api_manager +from autogpt.api_manager import ApiManager from autogpt.config import Config from autogpt.logs import logger from autogpt.types.openai import Message @@ -147,6 +147,7 @@ def create_chat_completion( ) if message is not None: return message + api_manager = ApiManager() response = None for attempt in range(num_retries): backoff = 2 ** (attempt + 2) @@ -228,6 +229,7 @@ def get_ada_embedding(text: str) -> List[float]: kwargs = {"model": model} embedding = create_embedding(text, **kwargs) + api_manager = ApiManager() api_manager.update_cost( prompt_tokens=embedding.usage.prompt_tokens, completion_tokens=0, diff --git a/autogpt/memory/base.py b/autogpt/memory/base.py index a6c92db6..d99cadde 100644 --- a/autogpt/memory/base.py +++ b/autogpt/memory/base.py @@ -1,11 +1,8 @@ """Base class for memory providers.""" import abc -from autogpt.config import Config from autogpt.singleton import AbstractSingleton -cfg = Config() - class MemoryProviderSingleton(AbstractSingleton): @abc.abstractmethod diff --git a/autogpt/prompts/prompt.py b/autogpt/prompts/prompt.py index cc06fabb..68d8a336 100644 --- a/autogpt/prompts/prompt.py +++ b/autogpt/prompts/prompt.py @@ -1,6 +1,6 @@ from colorama import Fore -from autogpt.api_manager import api_manager +from autogpt.api_manager import ApiManager from autogpt.config.ai_config import AIConfig from autogpt.config.config import Config from autogpt.logs import logger @@ -115,6 +115,7 @@ Continue (y/n): """ config.save(CFG.ai_settings_file) # set the total api budget + api_manager = ApiManager() api_manager.set_total_budget(config.api_budget) # Agent Created, print message diff --git a/tests/conftest.py b/tests/conftest.py index 9ab49189..745f8b5d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -3,7 +3,6 @@ from pathlib import Path import pytest from autogpt.api_manager import ApiManager -from autogpt.api_manager import api_manager as api_manager_ from autogpt.config import Config from autogpt.workspace import Workspace @@ -32,7 +31,6 @@ def config(workspace: Workspace) -> Config: @pytest.fixture() def api_manager() -> ApiManager: - old_attrs = api_manager_.__dict__.copy() - api_manager_.reset() - yield api_manager_ - api_manager_.__dict__.update(old_attrs) + if ApiManager in ApiManager._instances: + del ApiManager._instances[ApiManager] + return ApiManager() diff --git a/tests/unit/test_commands.py b/tests/unit/test_commands.py index a749cc6d..8596e1c9 100644 --- a/tests/unit/test_commands.py +++ b/tests/unit/test_commands.py @@ -3,8 +3,7 @@ from unittest.mock import MagicMock, patch import pytest -import autogpt.agent.agent_manager as agent_manager -from autogpt.app import execute_command, list_agents, start_agent +from autogpt.app import list_agents, start_agent from tests.utils import requires_api_key @@ -14,9 +13,11 @@ def test_make_agent() -> None: """Test that an agent can be created""" # Use the mock agent manager to avoid creating a real agent with patch("openai.ChatCompletion.create") as mock: - obj = MagicMock() - obj.response.choices[0].messages[0].content = "Test message" - mock.return_value = obj + response = MagicMock() + response.choices[0].messages[0].content = "Test message" + response.usage.prompt_tokens = 1 + response.usage.completion_tokens = 1 + mock.return_value = response start_agent("Test Agent", "chat", "Hello, how are you?", "gpt-3.5-turbo") agents = list_agents() assert "List of agents:\n0: chat" == agents