Export datasets to JSONL for fine-tuning

src/e2etests/judgment_client_test.py

@@ -12,6 +12,7 @@
 import string
 import logging
 from typing import Generator
+import json
 
 from judgeval.judgment_client import JudgmentClient
 from judgeval.data import Example
@@ -209,6 +210,46 @@ def test_assert_test(self, client: JudgmentClient):
                 override=True
             )
 
+    def test_export_jsonl(self, client: JudgmentClient, random_name: str):
+        """Test JSONL dataset export functionality."""
+        # Create and push test dataset
+        dataset = client.create_dataset()
+        dataset.add_example(Example(
+            input="Test input 1", 
+            actual_output="Test output 1",
+            expected_output="Expected output 1"
+        ))
+        dataset.add_ground_truth(GroundTruthExample(
+            input="GT input 1",
+            actual_output="GT output 1"
+        ))
+        client.push_dataset(alias=random_name, dataset=dataset, overwrite=True)
+
+        # Export as JSONL
+        response = client.eval_dataset_client.export_jsonl(random_name)
+        assert response.status_code == 200, "Export request failed"
+
+        # Validate JSONL format and content
+        example_count = 0
+        ground_truth_count = 0
+
+        for line in response.iter_lines():
+            if line:
+                entry = json.loads(line.decode('utf-8'))
+                assert "input" in entry, "Missing input field"
+                assert "output" in entry, "Missing output field"
+                assert "source" in entry, "Missing source field"
+
+                if entry["source"] == "example":
+                    example_count += 1
+                    assert "expected_output" in entry, "Example missing expected_output"
+                elif entry["source"] == "ground_truth":
+                    ground_truth_count += 1
+                    assert "source_file" not in entry, "Ground truth should not have source_file by default"
+
+        assert example_count == 1, f"Expected 1 example, got {example_count}"
+        assert ground_truth_count == 1, f"Expected 1 ground truth, got {ground_truth_count}"
+
 class TestAdvancedFeatures:
     def test_json_scorer(self, client: JudgmentClient):
         """Test JSON scorer functionality."""
@@ -537,7 +578,8 @@ def run_selected_tests(client, test_names: list[str]):
         'custom_judge_vertexai': test_custom_judges.test_custom_judge_vertexai,
         'fetch_traces': test_trace_operations.test_fetch_traces_by_time_period,
         'fetch_traces_invalid': test_trace_operations.test_fetch_traces_invalid_period,
-        'fetch_traces_missing_key': test_trace_operations.test_fetch_traces_missing_api_key
+        'fetch_traces_missing_key': test_trace_operations.test_fetch_traces_missing_api_key,
+        'export_jsonl': test_basic_operations.test_export_jsonl,
     }
 
     for test_name in test_names:
@@ -559,6 +601,7 @@ def run_selected_tests(client, test_names: list[str]):
 
     run_selected_tests(client, [
         'dataset',
+        'pull_all_user_dataset_stats',
         'run_eval', 
         'delete_eval_by_project_and_run_name',
         'delete_eval_by_project',
@@ -570,6 +613,6 @@ def run_selected_tests(client, test_names: list[str]):
         'custom_judge_vertexai',
         'fetch_traces',
         'fetch_traces_invalid',
-        'fetch_traces_missing_key'
-        'pull_all_user_dataset_stats',
+        'fetch_traces_missing_key',
+        'export_jsonl',
     ])

src/e2etests/test.py

@@ -0,0 +1,29 @@
+import requests
+from dotenv import load_dotenv
+import os
+
+def test():
+    response = requests.post(
+        "http://localhost:8000/organization/fetch_all_organizations_for_user/",
+        json={},
+        headers={"Content-Type": "application/json", "Authorization": f"Bearer {os.getenv('JUDGMENT_API_KEY')}"}
+    )
+    print(response.json())
+
+def test2():
+    response = requests.post(
+        "http://localhost:8000/organization/create/",
+        json={"organization_name": "test_organization2"},
+        headers={"Content-Type": "application/json", "Authorization": f"Bearer {os.getenv('JUDGMENT_API_KEY')}"}
+    )
+    print(response.json())
+
+
+
+def main():
+    load_dotenv()
+    test()
+
+
+if __name__ == "__main__":
+    main()

src/judgeval/constants.py

@@ -36,6 +36,7 @@ def _missing_(cls, value):
 JUDGMENT_EVAL_API_URL = f"{ROOT_API}/evaluate/"
 JUDGMENT_DATASETS_PUSH_API_URL = f"{ROOT_API}/datasets/push/"
 JUDGMENT_DATASETS_PULL_API_URL = f"{ROOT_API}/datasets/pull/"
+JUDGMENT_DATASETS_EXPORT_JSONL_API_URL = f"{ROOT_API}/datasets/export_jsonl/"
 JUDGMENT_DATASETS_PULL_ALL_API_URL = f"{ROOT_API}/datasets/get_all_stats/"
 JUDGMENT_DATASETS_EDIT_API_URL = f"{ROOT_API}/datasets/edit/"
 JUDGMENT_EVAL_LOG_API_URL = f"{ROOT_API}/log_eval_results/"

src/judgeval/data/datasets/eval_dataset_client.py

@@ -9,6 +9,7 @@
     JUDGMENT_DATASETS_PULL_API_URL, 
     JUDGMENT_DATASETS_PULL_ALL_API_URL,
     JUDGMENT_DATASETS_EDIT_API_URL,
+    JUDGMENT_DATASETS_EXPORT_JSONL_API_URL
 )
 from judgeval.data import Example
 from judgeval.data.datasets import EvalDataset
@@ -243,14 +244,47 @@ def edit_dataset(self, alias: str, examples: List[Example], ground_truths: List[
             except requests.exceptions.RequestException as e:
                 error(f"Error editing dataset: {str(e)}")
                 return False
+
+            info(f"Successfully edited dataset '{alias}'")
+            return True
 
+    def export_jsonl(self, alias: str) -> requests.Response:
+        """Export dataset in JSONL format from Judgment platform"""
+        debug(f"Exporting dataset with alias '{alias}' as JSONL")
+        with Progress(
+            SpinnerColumn(style="rgb(106,0,255)"),
+            TextColumn("[progress.description]{task.description}"),
+            transient=False,
+        ) as progress:
+            task_id = progress.add_task(
+                f"Exporting [rgb(106,0,255)]'{alias}'[/rgb(106,0,255)] as JSONL...",
+                total=100,
+            )
+            try:
+                response = requests.post(
+                    JUDGMENT_DATASETS_EXPORT_JSONL_API_URL,
+                    json={"alias": alias},
+                    headers={
+                        "Content-Type": "application/json",
+                        "Authorization": f"Bearer {self.judgment_api_key}"
+                    },
+                    stream=True
+                )
+                response.raise_for_status()
+            except requests.exceptions.HTTPError as err:
+                if err.response.status_code == 404:
+                    error(f"Dataset not found: {alias}")
+                else:
+                    error(f"HTTP error during export: {err}")
+                raise
+            except Exception as e:
+                error(f"Error during export: {str(e)}")
+                raise
+
+            info(f"Successfully exported dataset with alias '{alias}'")
             progress.update(
                 task_id,
                 description=f"{progress.tasks[task_id].description} [rgb(25,227,160)]Done!)",
             )
-
-            info(f"Successfully edited dataset '{alias}'")
-            return True
 
-
-
+            return response