aboutsummaryrefslogtreecommitdiff
path: root/tests/unit/test_api_manager.py
diff options
context:
space:
mode:
Diffstat (limited to 'tests/unit/test_api_manager.py')
-rw-r--r--tests/unit/test_api_manager.py86
1 files changed, 0 insertions, 86 deletions
diff --git a/tests/unit/test_api_manager.py b/tests/unit/test_api_manager.py
deleted file mode 100644
index 615204d19..000000000
--- a/tests/unit/test_api_manager.py
+++ /dev/null
@@ -1,86 +0,0 @@
-from unittest.mock import patch
-
-import pytest
-from pytest_mock import MockerFixture
-
-from autogpt.llm.api_manager import ApiManager
-from autogpt.llm.providers.openai import OPEN_AI_CHAT_MODELS, OPEN_AI_EMBEDDING_MODELS
-
-api_manager = ApiManager()
-
-
-@pytest.fixture(autouse=True)
-def reset_api_manager():
- api_manager.reset()
- yield
-
-
-@pytest.fixture(autouse=True)
-def mock_costs(mocker: MockerFixture):
- mocker.patch.multiple(
- OPEN_AI_CHAT_MODELS["gpt-3.5-turbo"],
- prompt_token_cost=0.0013,
- completion_token_cost=0.0025,
- )
- mocker.patch.multiple(
- OPEN_AI_EMBEDDING_MODELS["text-embedding-ada-002"],
- prompt_token_cost=0.0004,
- )
- yield
-
-
-class TestApiManager:
- def test_getter_methods(self):
- """Test the getter methods for total tokens, cost, and budget."""
- api_manager.update_cost(600, 1200, "gpt-3.5-turbo")
- api_manager.set_total_budget(10.0)
- assert api_manager.get_total_prompt_tokens() == 600
- assert api_manager.get_total_completion_tokens() == 1200
- assert api_manager.get_total_cost() == (600 * 0.0013 + 1200 * 0.0025) / 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_completion_model():
- """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() == prompt_tokens
- assert api_manager.get_total_completion_tokens() == completion_tokens
- assert (
- api_manager.get_total_cost()
- == (prompt_tokens * 0.0013 + completion_tokens * 0.0025) / 1000
- )
-
- @staticmethod
- def test_update_cost_embedding_model():
- """Test if updating the cost works correctly."""
- prompt_tokens = 1337
- model = "text-embedding-ada-002"
-
- api_manager.update_cost(prompt_tokens, 0, model)
-
- assert api_manager.get_total_prompt_tokens() == prompt_tokens
- assert api_manager.get_total_completion_tokens() == 0
- assert api_manager.get_total_cost() == (prompt_tokens * 0.0004) / 1000
-
- @staticmethod
- def test_get_models():
- """Test if getting models works correctly."""
- with patch("openai.Model.list") as mock_list_models:
- mock_list_models.return_value = {"data": [{"id": "gpt-3.5-turbo"}]}
- result = api_manager.get_models()
-
- assert result[0]["id"] == "gpt-3.5-turbo"
- assert api_manager.models[0]["id"] == "gpt-3.5-turbo"