Fix custom_python not being copied (#5512)

2 files changed, 12 insertions, 6 deletions

diff --git a/benchmark/agbenchmark/app.py b/benchmark/agbenchmark/app.py
index d78c19ff8..ad14cb692 100644
--- a/benchmark/agbenchmark/app.py
+++ b/benchmark/agbenchmark/app.py
@@ -331,15 +331,22 @@ async def proxy(request: Request, task_id: str):
 
 @router.post("/agent/tasks/{task_id}/evaluations")
 async def create_evaluation(task_id: str) -> deque:
+    from agbenchmark.__main__ import TEMP_FOLDER_ABS_PATH
     from agbenchmark.agent_api_interface import copy_agent_artifacts_into_temp_folder
+    from agbenchmark.agent_interface import copy_artifacts_into_temp_folder
     from agbenchmark.generate_test import create_challenge
 
     try:
         async with ApiClient(configuration) as api_client:
             api_instance = AgentApi(api_client)
             await copy_agent_artifacts_into_temp_folder(api_instance, task_id)
-
+        # add custom python
         data = CHALLENGES[task_informations[task_id]["eval_id"]]
+
+        artifact_path = str(Path(data["path"]).parent)
+        copy_artifacts_into_temp_folder(
+            TEMP_FOLDER_ABS_PATH, "custom_python", artifact_path
+        )
         json_file = CHALLENGES[task_informations[task_id]["eval_id"]]["path"]
         json_files = deque()
 
diff --git a/benchmark/agbenchmark/utils/challenge.py b/benchmark/agbenchmark/utils/challenge.py
index 49d067664..36bf0af22 100644
--- a/benchmark/agbenchmark/utils/challenge.py
+++ b/benchmark/agbenchmark/utils/challenge.py
@@ -49,11 +49,6 @@ class Challenge(ABC):
     async def setup_challenge(self, config: Dict[str, Any], cutoff: int) -> None:
         from agbenchmark.agent_interface import copy_artifacts_into_temp_folder
 
-        artifact_paths = [
-            self.ARTIFACTS_LOCATION,
-            str(Path(self.CHALLENGE_LOCATION).parent),
-        ]
-
         if not self.task:
             return
 
@@ -66,6 +61,10 @@ class Challenge(ABC):
 
         # hidden files are added after the agent runs. Hidden files can be python test files.
         # We copy them in the temporary folder to make it easy to import the code produced by the agent
+        artifact_paths = [
+            self.ARTIFACTS_LOCATION,
+            str(Path(self.CHALLENGE_LOCATION).parent),
+        ]
         for path in artifact_paths:
             copy_artifacts_into_temp_folder(TEMP_FOLDER_ABS_PATH, "custom_python", path)