Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Have api manager use singleton pattern #3269

Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 10 additions & 15 deletions autogpt/api_manager.py
Original file line number Diff line number Diff line change
@@ -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
from autogpt.singleton import Singleton

cfg = Config()
print_total_cost = cfg.debug_mode


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
ntindle marked this conversation as resolved.
Show resolved Hide resolved

def reset(self):
self.total_prompt_tokens = 0
Expand All @@ -28,7 +25,7 @@ def create_chat_completion(
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:
Expand All @@ -42,6 +39,9 @@ def create_chat_completion(
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,
Expand All @@ -59,8 +59,7 @@ def create_chat_completion(
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)
Expand All @@ -81,8 +80,7 @@ def update_cost(self, prompt_tokens, completion_tokens, model):
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):
"""
Expand Down Expand Up @@ -128,6 +126,3 @@ def get_total_budget(self):
float: The total budget for API calls.
"""
return self.total_budget


api_manager = ApiManager(cfg.debug_mode)
3 changes: 2 additions & 1 deletion autogpt/chat.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
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
Expand Down Expand Up @@ -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 = (
Expand Down
4 changes: 3 additions & 1 deletion autogpt/llm_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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,
Expand Down
3 changes: 0 additions & 3 deletions autogpt/memory/base.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down
3 changes: 2 additions & 1 deletion autogpt/prompts/prompt.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -115,6 +115,7 @@ def construct_main_ai_config() -> AIConfig:
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
Expand Down
8 changes: 3 additions & 5 deletions tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
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

Expand Down Expand Up @@ -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()
11 changes: 6 additions & 5 deletions tests/unit/test_commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,7 @@

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


Expand All @@ -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
Expand Down