### autogenerated_evaluation_dataset_creation_example_script.py ###
# This script demonstrates how to create a autogenerated evaluation dataset
# For autogenerated evaluation datasets, a generation job workflow is created to generate test cases.
# Test cases are generated based on the knowledge base provided, and must be approved before the dataset can be published.
# Evaluation datasets, once published, can be used for application variant runs and report card generation
# For the purposes of this script, we assume that you already have a knowledge base set up with artifacts uploaded
# This can be either done through API, or by going to https://egp.dashboard.scale.com/knowledge-bases/new
import os
import time

from scale_gp import SGPClient

account_id = os.environ.get("SGP_ACCOUNT_ID", None)
api_key = os.environ.get("SGP_API_KEY", None)
knowledge_base_id = os.environ.get("SGP_KNOWLEDGE_BASE_ID", None)

assert (
    account_id is not None
), "You need to set the SGP_ACCOUNT_ID - you can find it at https://egp.dashboard.scale.com/admin/accounts"
assert api_key is not None, "You need to provide your API key - see https://egp.dashboard.scale.com/admin/api-key"

# Create an SGP Client
client = SGPClient(api_key=api_key)

### autogenerated evaluation dataset ###

# we assume we have a knowledge base with artifacts already uploaded
autogenerated_evaluation_dataset = client.evaluation_datasets.create(
    account_id=account_id,
    name="autogenerated_eval_dataset",
    schema_type="GENERATION",
    type="autogenerated",
    knowledge_base_id=knowledge_base_id,
)

# view evaluation dataset details. This will show the dataset_id, name, schema_type, type, harms_list, advanced_config etc.
dataset = client.evaluation_datasets.retrieve(evaluation_dataset_id=autogenerated_evaluation_dataset.id)

# start generation job to generate test cases
generation_job = client.evaluation_datasets.generation_jobs.create(
    evaluation_dataset_id=autogenerated_evaluation_dataset.id,
    num_test_cases=3,
    group_by_artifact_id=False,
)

while True:
    if generation_job.status == "Pending":
        print("generating test cases...")
        time.sleep(5)
    else:
        break

# view autogenerated test cases
test_cases = client.evaluation_datasets.autogenerated_draft_test_cases.list(
    evaluation_dataset_id=autogenerated_evaluation_dataset.id
)

# approve autogenerated test cases. all test cases must be approved before the dataset can be published.
for test_case in test_cases.items:
    client.evaluation_datasets.autogenerated_draft_test_cases.approve(
        evaluation_dataset_id=autogenerated_evaluation_dataset.id,
        autogenerated_draft_test_case_id=test_case.id,
    )

# snapshot into a dataset version when test cases are approved
autogenerated_dataset_version = client.evaluation_datasets.evaluation_dataset_versions.create(
    evaluation_dataset_id=autogenerated_evaluation_dataset.id
)

# publish the dataset. datasets must be published before they can be used in evaluations.
published_dataset_response = client.evaluation_datasets.publish(
    evaluation_dataset_id=autogenerated_evaluation_dataset.id,
)

print("Published autogenerated evaluation dataset with id:", autogenerated_evaluation_dataset.id)
print(autogenerated_dataset_version)

test_cases = client.evaluation_datasets.test_cases.list(
    evaluation_dataset_id=autogenerated_evaluation_dataset.id, account_id=account_id
)

print("Created the following test cases:")
print(test_cases.items)
