Fix send to gdrive and tracking the wrong challenge name (#225)

Signed-off-by: Merwane Hamadi <merwanehamadi@gmail.com>

agbenchmark/conftest.py

@@ -8,7 +8,6 @@ from pathlib import Path  # noqa
 from typing import Any, Dict, Generator
 
 import pytest
-from helicone.lock import HeliconeLockManager
 
 from agbenchmark.reports.reports import (
     finalize_reports,
@@ -169,7 +168,6 @@ def pytest_runtest_makereport(item: Any, call: Any) -> None:
             test_name = challenge_data["name"]
             generate_single_call_report(item, call, challenge_data)
         # else: it's a same_task=false suite (tests aren't combined)
-        HeliconeLockManager.write_custom_property("challenge", test_name)
     if call.when == "teardown":
         finalize_reports(item, challenge_data)
 

agbenchmark/generate_test.py

@@ -56,6 +56,9 @@ def create_single_test(
 
     # Define test method within the dynamically created class
     def test_method(self, config: Dict[str, Any], request) -> None:  # type: ignore
+        from helicone.lock import HeliconeLockManager
+        HeliconeLockManager.write_custom_property("challenge", self.data.name)
+
         cutoff = self.data.cutoff or 60
         self.setup_challenge(config, cutoff)
 

agbenchmark/utils/get_data_from_helicone.py

@@ -79,3 +79,4 @@ query ExampleQuery($properties: [PropertyFilter!]){
         )
     except Exception as err:
         print(f"Error occurred: {err}")
+        raise

get_data_from_helicone.py

@@ -1,45 +0,0 @@
-import json
-
-import requests
-
-# Define the endpoint of your GraphQL server
-url = "https://www.helicone.ai/api/graphql"
-
-# Set the headers, usually you'd need to set the content type and possibly an authorization token
-headers = {"authorization": "Bearer sk-"}
-
-# Define the query, variables, and operation name
-query = """
-query ExampleQuery($limit: Int, $filters: [HeliconeRequestFilter!]) {
-  user {
-    id
-  }
-  heliconeRequest(limit: $limit,filters: $filters) {
-    responseBody
-  }
-}
-"""
-
-variables = {
-    "limit": 100,
-    "filters": [{"property": {"value": {"equals": "beebot"}, "name": "agent"}}],
-}
-
-operation_name = "ExampleQuery"
-
-# Make the request
-response = requests.post(
-    url,
-    headers=headers,
-    json={"query": query, "variables": variables, "operationName": operation_name},
-)
-data = response.json()
-total_tokens_sum = 0
-
-for item in data["data"]["heliconeRequest"]:
-    total_tokens_sum += item["responseBody"]["usage"]["total_tokens"]
-
-# Extract the data from the response (consider adding error checks)
-
-print(json.dumps(data, indent=4, ensure_ascii=False))
-print(total_tokens_sum)

send_to_googledrive.py

@@ -71,26 +71,28 @@ def process_test(
 
 
 # Loop over each directory in the base directory
-for sub_dir in os.listdir(base_dir):
-    # Define the subdirectory path
-    sub_dir_path = os.path.join(base_dir, sub_dir)
+for agent_dir in os.listdir(base_dir):
+    agent_dir_path = os.path.join(base_dir, agent_dir)
 
-    # Ensure the sub_dir_path is a directory
-    if os.path.isdir(sub_dir_path):
-        # Loop over each file in the subdirectory
-        for filename in os.listdir(sub_dir_path):
-            # Check if the file is a JSON file
-            if filename.endswith(".json"):
-                # Define the file path
-                file_path = os.path.join(sub_dir_path, filename)
+    # Ensure the agent_dir_path is a directory
+    if os.path.isdir(agent_dir_path):
+        # Loop over each sub-directory in the agent directory (e.g., "folder49_07-28-03-53")
+        for report_folder in os.listdir(agent_dir_path):
+            report_folder_path = os.path.join(agent_dir_path, report_folder)
 
-                # Load the JSON data from the file
-                with open(file_path, "r") as f:
-                    data = json.load(f)
+            # Ensure the report_folder_path is a directory
+            if os.path.isdir(report_folder_path):
+                # Check for a file named "report.json" in the sub-directory
+                report_path = os.path.join(report_folder_path, "report.json")
 
-                # Loop through each test
-                for test_name, test_info in data["tests"].items():
-                    process_test(test_name, test_info, sub_dir, data)
+                if os.path.exists(report_path):
+                    # Load the JSON data from the file
+                    with open(report_path, "r") as f:
+                        data = json.load(f)
+
+                    # Loop through each test
+                    for test_name, test_info in data["tests"].items():
+                        process_test(test_name, test_info, agent_dir, data)
 
 # Convert the list of rows into a DataFrame
 df = pd.DataFrame(rows)