AutoGPT

1
import pytest
2
from pytest_mock import MockerFixture
3

4
from autogpt.core.resource.model_providers import (
5
    OPEN_AI_CHAT_MODELS,
6
    OPEN_AI_EMBEDDING_MODELS,
7
)
8
from autogpt.llm.api_manager import ApiManager
9

10
api_manager = ApiManager()
11

12

13
@pytest.fixture(autouse=True)
14
def reset_api_manager():
15
    api_manager.reset()
16
    yield
17

18

19
@pytest.fixture(autouse=True)
20
def mock_costs(mocker: MockerFixture):
21
    mocker.patch.multiple(
22
        OPEN_AI_CHAT_MODELS["gpt-3.5-turbo"],
23
        prompt_token_cost=0.0013,
24
        completion_token_cost=0.0025,
25
    )
26
    mocker.patch.multiple(
27
        OPEN_AI_EMBEDDING_MODELS["text-embedding-ada-002"],
28
        prompt_token_cost=0.0004,
29
    )
30
    yield
31

32

33
class TestApiManager:
34
    def test_getter_methods(self):
35
        """Test the getter methods for total tokens, cost, and budget."""
36
        api_manager.update_cost(600, 1200, "gpt-3.5-turbo")
37
        api_manager.set_total_budget(10.0)
38
        assert api_manager.get_total_prompt_tokens() == 600
39
        assert api_manager.get_total_completion_tokens() == 1200
40
        assert api_manager.get_total_cost() == (600 * 0.0013 + 1200 * 0.0025) / 1000
41
        assert api_manager.get_total_budget() == 10.0
42

43
    @staticmethod
44
    def test_set_total_budget():
45
        """Test if setting the total budget works correctly."""
46
        total_budget = 10.0
47
        api_manager.set_total_budget(total_budget)
48

49
        assert api_manager.get_total_budget() == total_budget
50

51
    @staticmethod
52
    def test_update_cost_completion_model():
53
        """Test if updating the cost works correctly."""
54
        prompt_tokens = 50
55
        completion_tokens = 100
56
        model = "gpt-3.5-turbo"
57

58
        api_manager.update_cost(prompt_tokens, completion_tokens, model)
59

60
        assert api_manager.get_total_prompt_tokens() == prompt_tokens
61
        assert api_manager.get_total_completion_tokens() == completion_tokens
62
        assert (
63
            api_manager.get_total_cost()
64
            == (prompt_tokens * 0.0013 + completion_tokens * 0.0025) / 1000
65
        )
66

67
    @staticmethod
68
    def test_update_cost_embedding_model():
69
        """Test if updating the cost works correctly."""
70
        prompt_tokens = 1337
71
        model = "text-embedding-ada-002"
72

73
        api_manager.update_cost(prompt_tokens, 0, model)
74

75
        assert api_manager.get_total_prompt_tokens() == prompt_tokens
76
        assert api_manager.get_total_completion_tokens() == 0
77
        assert api_manager.get_total_cost() == (prompt_tokens * 0.0004) / 1000
78