From 3fc1a9572b46de5a06516c3be02ec9a4ef7a1b33 Mon Sep 17 00:00:00 2001
From: Merwane Hamadi <merwanehamadi@gmail.com>
Date: Fri, 14 Jul 2023 08:35:02 -0700
Subject: [PATCH] Add basic code generation challenge

Signed-off-by: Merwane Hamadi <merwanehamadi@gmail.com>
---
 agbenchmark/challenge.py                      |  7 +++++
 .../code/d4/artifacts_out/__init__.py         |  0
 .../challenges/code/d4/artifacts_out/code.py  | 12 +++++++
 agbenchmark/challenges/code/d4/data.json      | 18 +++++++++++
 .../challenges/code/d4/hidden_files/test.py   | 31 +++++++++++++++++++
 agent/gpt-engineer                            |  2 +-
 6 files changed, 69 insertions(+), 1 deletion(-)
 create mode 100644 agbenchmark/challenges/code/d4/artifacts_out/__init__.py
 create mode 100644 agbenchmark/challenges/code/d4/artifacts_out/code.py
 create mode 100644 agbenchmark/challenges/code/d4/data.json
 create mode 100644 agbenchmark/challenges/code/d4/hidden_files/test.py

diff --git a/agbenchmark/challenge.py b/agbenchmark/challenge.py
index aeebd7ad897..874fd45bd99 100644
--- a/agbenchmark/challenge.py
+++ b/agbenchmark/challenge.py
@@ -46,6 +46,13 @@ def setup_challenge(self, config: Dict[str, Any]) -> None:
 
         run_agent(self.task, config, self.CHALLENGE_LOCATION)
 
+        # hidden files are added after the agent runs. Hidden files can be python test files.
+        # We copy them in the workspace to make it easy to import the code produced by the agent
+
+        copy_artifacts_into_workspace(
+            config["workspace"], "hidden_files", self.CHALLENGE_LOCATION
+        )
+
     def test_method(self, config: Dict[str, Any]) -> None:
         raise NotImplementedError
 
diff --git a/agbenchmark/challenges/code/d4/artifacts_out/__init__.py b/agbenchmark/challenges/code/d4/artifacts_out/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/agbenchmark/challenges/code/d4/artifacts_out/code.py b/agbenchmark/challenges/code/d4/artifacts_out/code.py
new file mode 100644
index 00000000000..de3d8c62cad
--- /dev/null
+++ b/agbenchmark/challenges/code/d4/artifacts_out/code.py
@@ -0,0 +1,12 @@
+# mypy: ignore-errors
+from typing import List, Optional
+
+
+def two_sum(nums: List, target: int) -> Optional[List[int]]:
+    seen = {}
+    for i, num in enumerate(nums):
+        complement = target - num
+        if complement in seen:
+            return [seen[complement], i]
+        seen[num] = i
+    return None
diff --git a/agbenchmark/challenges/code/d4/data.json b/agbenchmark/challenges/code/d4/data.json
new file mode 100644
index 00000000000..224fbc80414
--- /dev/null
+++ b/agbenchmark/challenges/code/d4/data.json
@@ -0,0 +1,18 @@
+{
+  "name": "TestBasicCodeGeneration",
+  "category": ["code"],
+  "task": "Create a two_sum function in a file called code.py. Given an array of integers, return indices of the two numbers such that they add up to a specific target. You may assume that each input would have exactly one solution, and you may not use the same element twice. Example: Given nums = [2, 7, 11, 15], target = 9, Because nums[0] + nums[1] = 2 + 7 = 9, return [0, 1].",
+  "dependencies": ["TestWriteFile"],
+  "ground": {
+    "answer": "[0, 1] [2, 5] [0, 3]",
+    "should_contain": ["[0, 1]", "[2, 5]", "[0, 3]"],
+    "should_not_contain": [],
+    "files": ["test.py"],
+    "type": "execute_python_code"
+  },
+  "info": {
+    "difficulty": "basic",
+    "description": "Tests ability for the agent to create the two_sum function.",
+    "side_effects": []
+  }
+}
diff --git a/agbenchmark/challenges/code/d4/hidden_files/test.py b/agbenchmark/challenges/code/d4/hidden_files/test.py
new file mode 100644
index 00000000000..d85d1353758
--- /dev/null
+++ b/agbenchmark/challenges/code/d4/hidden_files/test.py
@@ -0,0 +1,31 @@
+# mypy: ignore-errors
+from code import two_sum
+from typing import List
+
+
+def test_two_sum(nums: List, target: int, expected_result: List[int]) -> None:
+    result = two_sum(nums, target)
+    print(result)
+    assert (
+        result == expected_result
+    ), f"AssertionError: Expected the output to be {expected_result}"
+
+
+if __name__ == "__main__":
+    # test the trivial case with the first two numbers
+    nums = [2, 7, 11, 15]
+    target = 9
+    expected_result = [0, 1]
+    test_two_sum(nums, target, expected_result)
+
+    # test for ability to use zero and the same number twice
+    nums = [2, 7, 0, 15, 12, 0]
+    target = 0
+    expected_result = [2, 5]
+    test_two_sum(nums, target, expected_result)
+
+    # test for first and last index usage and negative numbers
+    nums = [-6, 7, 11, 4]
+    target = -2
+    expected_result = [0, 3]
+    test_two_sum(nums, target, expected_result)
diff --git a/agent/gpt-engineer b/agent/gpt-engineer
index 521d626c007..bca191cd76c 160000
--- a/agent/gpt-engineer
+++ b/agent/gpt-engineer
@@ -1 +1 @@
-Subproject commit 521d626c0075ed6545f01b771757c856f8addbd6
+Subproject commit bca191cd76cdea0335da91d004c64d9bb8520fea