diff --git a/gpt_engineer/db.py b/gpt_engineer/db.py index 0d7e64e..1f117b2 100644 --- a/gpt_engineer/db.py +++ b/gpt_engineer/db.py @@ -1,5 +1,6 @@ from dataclasses import dataclass from pathlib import Path +from typing import Optional # This class represents a simple database that stores its data as files in a directory. @@ -47,3 +48,4 @@ class DBs: preprompts: DB input: DB workspace: DB + archive: DB diff --git a/gpt_engineer/main.py b/gpt_engineer/main.py index 9dea702..81c4c4c 100644 --- a/gpt_engineer/main.py +++ b/gpt_engineer/main.py @@ -1,6 +1,6 @@ import json import logging -import shutil +import os from pathlib import Path @@ -18,7 +18,6 @@ app = typer.Typer() @app.command() def main( project_path: str = typer.Argument("example", help="path"), - delete_existing: bool = typer.Argument(False, help="delete existing files"), model: str = typer.Argument("gpt-4", help="model id string"), temperature: float = 0.1, steps_config: steps.Config = typer.Option( @@ -35,28 +34,25 @@ def main( ): logging.basicConfig(level=logging.DEBUG if verbose else logging.INFO) - input_path = Path(project_path).absolute() - memory_path = input_path / f"{run_prefix}memory" - workspace_path = input_path / f"{run_prefix}workspace" - - if delete_existing: - # Delete files and subdirectories in paths - shutil.rmtree(memory_path, ignore_errors=True) - shutil.rmtree(workspace_path, ignore_errors=True) - model = fallback_model(model) - ai = AI( model=model, temperature=temperature, ) + input_path = Path(project_path).absolute() + memory_path = input_path / f"{run_prefix}memory" + workspace_path = input_path / f"{run_prefix}workspace" + archive_path = input_path / f"{run_prefix}archive" + + initial_run = not os.path.exists(memory_path) and not os.path.exists(workspace_path) dbs = DBs( memory=DB(memory_path), logs=DB(memory_path / "logs"), input=DB(input_path), workspace=DB(workspace_path), preprompts=DB(Path(__file__).parent / "preprompts"), + archive=None if initial_run else DB(archive_path), ) steps = STEPS[steps_config] diff --git a/gpt_engineer/steps.py b/gpt_engineer/steps.py index 7a012f4..438b162 100644 --- a/gpt_engineer/steps.py +++ b/gpt_engineer/steps.py @@ -1,5 +1,8 @@ +import datetime import json +import os import re +import shutil import subprocess from enum import Enum @@ -258,6 +261,16 @@ def fix_code(ai: AI, dbs: DBs): return messages +def archive(ai: AI, dbs: DBs): + os.makedirs(dbs.archive.path, exist_ok=True) + timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + shutil.move(dbs.memory.path, dbs.archive.path / timestamp / dbs.memory.path.name) + shutil.move( + dbs.workspace.path, dbs.archive.path / timestamp / dbs.workspace.path.name + ) + return [] + + def human_review(ai: AI, dbs: DBs): review = human_input() dbs.memory["review"] = review.to_json() # type: ignore @@ -280,15 +293,17 @@ class Config(str, Enum): # Different configs of what steps to run STEPS = { Config.DEFAULT: [ + archive, clarify, gen_clarified_code, gen_entrypoint, execute_entrypoint, human_review, ], - Config.BENCHMARK: [simple_gen, gen_entrypoint], - Config.SIMPLE: [simple_gen, gen_entrypoint, execute_entrypoint], + Config.BENCHMARK: [archive, simple_gen, gen_entrypoint], + Config.SIMPLE: [archive, simple_gen, gen_entrypoint, execute_entrypoint], Config.TDD: [ + archive, gen_spec, gen_unit_tests, gen_code, @@ -297,6 +312,7 @@ STEPS = { human_review, ], Config.TDD_PLUS: [ + archive, gen_spec, gen_unit_tests, gen_code, @@ -306,6 +322,7 @@ STEPS = { human_review, ], Config.CLARIFY: [ + archive, clarify, gen_clarified_code, gen_entrypoint, @@ -313,6 +330,7 @@ STEPS = { human_review, ], Config.RESPEC: [ + archive, gen_spec, respec, gen_unit_tests, diff --git a/tests/steps/__init__.py b/tests/steps/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/steps/test_archive.py b/tests/steps/test_archive.py new file mode 100644 index 0000000..3211d3d --- /dev/null +++ b/tests/steps/test_archive.py @@ -0,0 +1,44 @@ +import datetime +import os + +from unittest.mock import MagicMock + +from gpt_engineer.db import DB, DBs +from gpt_engineer.steps import archive + + +def freeze_at(monkeypatch, time): + datetime_mock = MagicMock(wraps=datetime.datetime) + datetime_mock.now.return_value = time + monkeypatch.setattr(datetime, "datetime", datetime_mock) + + +def setup_dbs(tmp_path, dir_names): + directories = [tmp_path / name for name in dir_names] + + # Create DB objects + dbs = [DB(dir) for dir in directories] + + # Create DBs instance + return DBs(*dbs) + + +def test_archive(tmp_path, monkeypatch): + dbs = setup_dbs( + tmp_path, ["memory", "logs", "preprompts", "input", "workspace", "archive"] + ) + freeze_at(monkeypatch, datetime.datetime(2020, 12, 25, 17, 5, 55)) + archive(None, dbs) + assert not os.path.exists(tmp_path / "memory") + assert not os.path.exists(tmp_path / "workspace") + assert os.path.isdir(tmp_path / "archive" / "20201225_170555") + + dbs = setup_dbs( + tmp_path, ["memory", "logs", "preprompts", "input", "workspace", "archive"] + ) + freeze_at(monkeypatch, datetime.datetime(2022, 8, 14, 8, 5, 12)) + archive(None, dbs) + assert not os.path.exists(tmp_path / "memory") + assert not os.path.exists(tmp_path / "workspace") + assert os.path.isdir(tmp_path / "archive" / "20201225_170555") + assert os.path.isdir(tmp_path / "archive" / "20220814_080512") diff --git a/tests/test_collect.py b/tests/test_collect.py index e75cbc1..c02ed8b 100644 --- a/tests/test_collect.py +++ b/tests/test_collect.py @@ -19,7 +19,7 @@ def test_collect_learnings(monkeypatch): model = "test_model" temperature = 0.5 steps = [gen_code] - dbs = DBs(DB("/tmp"), DB("/tmp"), DB("/tmp"), DB("/tmp"), DB("/tmp")) + dbs = DBs(DB("/tmp"), DB("/tmp"), DB("/tmp"), DB("/tmp"), DB("/tmp"), DB("/tmp")) dbs.input = { "prompt": "test prompt\n with newlines", "feedback": "test feedback", diff --git a/tests/test_db.py b/tests/test_db.py index 072c3fa..e490aa3 100644 --- a/tests/test_db.py +++ b/tests/test_db.py @@ -27,7 +27,7 @@ def test_DB_operations(tmp_path): def test_DBs_initialization(tmp_path): - dir_names = ["memory", "logs", "preprompts", "input", "workspace"] + dir_names = ["memory", "logs", "preprompts", "input", "workspace", "archive"] directories = [tmp_path / name for name in dir_names] # Create DB objects @@ -41,6 +41,7 @@ def test_DBs_initialization(tmp_path): assert isinstance(dbs_instance.preprompts, DB) assert isinstance(dbs_instance.input, DB) assert isinstance(dbs_instance.workspace, DB) + assert isinstance(dbs_instance.archive, DB) def test_invalid_path(): @@ -106,11 +107,11 @@ def test_DBs_instantiation_with_wrong_number_of_arguments(tmp_path): DBs(db, db, db) with pytest.raises(TypeError): - DBs(db, db, db, db, db, db) + DBs(db, db, db, db, db, db, db) def test_DBs_dataclass_attributes(tmp_path): - dir_names = ["memory", "logs", "preprompts", "input", "workspace"] + dir_names = ["memory", "logs", "preprompts", "input", "workspace", "archive"] directories = [tmp_path / name for name in dir_names] # Create DB objects