AutoGPT

1
from __future__ import annotations
2

3
import os
4
import uuid
5
from pathlib import Path
6
from tempfile import TemporaryDirectory
7

8
import pytest
9
import yaml
10
from pytest_mock import MockerFixture
11

12
from autogpt.agents.agent import Agent, AgentConfiguration, AgentSettings
13
from autogpt.app.main import _configure_openai_provider
14
from autogpt.config import AIProfile, Config, ConfigBuilder
15
from autogpt.core.resource.model_providers import ChatModelProvider, OpenAIProvider
16
from autogpt.file_storage.local import (
17
    FileStorage,
18
    FileStorageConfiguration,
19
    LocalFileStorage,
20
)
21
from autogpt.llm.api_manager import ApiManager
22
from autogpt.logs.config import configure_logging
23
from autogpt.models.command_registry import CommandRegistry
24

25
pytest_plugins = [
26
    "tests.integration.agent_factory",
27
    "tests.integration.memory.utils",
28
    "tests.vcr",
29
]
30

31

32
@pytest.fixture()
33
def tmp_project_root(tmp_path: Path) -> Path:
34
    return tmp_path
35

36

37
@pytest.fixture()
38
def app_data_dir(tmp_project_root: Path) -> Path:
39
    dir = tmp_project_root / "data"
40
    dir.mkdir(parents=True, exist_ok=True)
41
    return dir
42

43

44
@pytest.fixture()
45
def storage(app_data_dir: Path) -> FileStorage:
46
    storage = LocalFileStorage(
47
        FileStorageConfiguration(root=app_data_dir, restrict_to_root=False)
48
    )
49
    storage.initialize()
50
    return storage
51

52

53
@pytest.fixture
54
def temp_plugins_config_file():
55
    """
56
    Create a plugins_config.yaml file in a temp directory
57
    so that it doesn't mess with existing ones.
58
    """
59
    config_directory = TemporaryDirectory()
60
    config_file = Path(config_directory.name) / "plugins_config.yaml"
61
    with open(config_file, "w+") as f:
62
        f.write(yaml.dump({}))
63

64
    yield config_file
65

66

67
@pytest.fixture(scope="function")
68
def config(
69
    temp_plugins_config_file: Path,
70
    tmp_project_root: Path,
71
    app_data_dir: Path,
72
    mocker: MockerFixture,
73
):
74
    if not os.environ.get("OPENAI_API_KEY"):
75
        os.environ["OPENAI_API_KEY"] = "sk-dummy"
76
    config = ConfigBuilder.build_config_from_env(project_root=tmp_project_root)
77

78
    config.app_data_dir = app_data_dir
79

80
    config.plugins_dir = "tests/unit/data/test_plugins"
81
    config.plugins_config_file = temp_plugins_config_file
82

83
    config.logging.log_dir = Path(__file__).parent / "logs"
84
    config.logging.plain_console_output = True
85
    config.noninteractive_mode = True
86

87
    # avoid circular dependency
88
    from autogpt.plugins.plugins_config import PluginsConfig
89

90
    config.plugins_config = PluginsConfig.load_config(
91
        plugins_config_file=config.plugins_config_file,
92
        plugins_denylist=config.plugins_denylist,
93
        plugins_allowlist=config.plugins_allowlist,
94
    )
95
    yield config
96

97

98
@pytest.fixture(scope="session")
99
def setup_logger(config: Config):
100
    configure_logging(**config.logging.dict())
101

102

103
@pytest.fixture()
104
def api_manager() -> ApiManager:
105
    if ApiManager in ApiManager._instances:
106
        del ApiManager._instances[ApiManager]
107
    return ApiManager()
108

109

110
@pytest.fixture
111
def llm_provider(config: Config) -> OpenAIProvider:
112
    return _configure_openai_provider(config)
113

114

115
@pytest.fixture
116
def agent(
117
    config: Config, llm_provider: ChatModelProvider, storage: FileStorage
118
) -> Agent:
119
    ai_profile = AIProfile(
120
        ai_name="Base",
121
        ai_role="A base AI",
122
        ai_goals=[],
123
    )
124

125
    command_registry = CommandRegistry()
126

127
    agent_prompt_config = Agent.default_settings.prompt_config.copy(deep=True)
128
    agent_prompt_config.use_functions_api = config.openai_functions
129

130
    agent_settings = AgentSettings(
131
        name=Agent.default_settings.name,
132
        description=Agent.default_settings.description,
133
        agent_id=f"AutoGPT-test-agent-{str(uuid.uuid4())[:8]}",
134
        ai_profile=ai_profile,
135
        config=AgentConfiguration(
136
            fast_llm=config.fast_llm,
137
            smart_llm=config.smart_llm,
138
            allow_fs_access=not config.restrict_to_workspace,
139
            use_functions_api=config.openai_functions,
140
            plugins=config.plugins,
141
        ),
142
        prompt_config=agent_prompt_config,
143
        history=Agent.default_settings.history.copy(deep=True),
144
    )
145

146
    agent = Agent(
147
        settings=agent_settings,
148
        llm_provider=llm_provider,
149
        command_registry=command_registry,
150
        file_storage=storage,
151
        legacy_config=config,
152
    )
153
    return agent
154