GenerateTestSetWithFilesExample.java

package ai.rhesis.sdk.examples;

import ai.rhesis.sdk.RhesisClient;
import ai.rhesis.sdk.entities.Test;
import ai.rhesis.sdk.entities.TestSet;
import ai.rhesis.sdk.synthesizers.GenerationConfig;
import ai.rhesis.sdk.synthesizers.MultiTurnSynthesizer;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.Arrays;
import java.util.List;

public class GenerateTestSetWithFilesExample {
  public static void main(String[] args) throws Exception {
    // Initialize the global default client
    RhesisClient client = RhesisClient.builder().apiKey(System.getenv("RHESIS_API_KEY")).build();
    RhesisClient.setDefault(client);

    // Configure the synthesizer
    System.out.println("Initializing MultiTurnSynthesizer...");
    GenerationConfig config =
        GenerationConfig.builder()
            .generationPrompt(
                "You are testing an HR document processor. Generate tests involving reading policy documents.")
            .testSetName("HR Document Processor Tests")
            .requirements(
                Arrays.asList("Accurately summarizes policies", "Identifies vacation days"))
            .categories(Arrays.asList("Functionality", "Document Processing"))
            .topics(Arrays.asList("HR", "Time Off"))
            .build();

    MultiTurnSynthesizer synthesizer = new MultiTurnSynthesizer(config, 5);

    // Generate the tests
    System.out.println("Generating test set (this may take a moment)...");
    TestSet generatedTestSet = synthesizer.generate(2);

    System.out.println("Generated TestSet locally with name: " + generatedTestSet.name());
    System.out.println("Number of tests generated: " + generatedTestSet.tests().size());

    // Push the test set to the Rhesis platform
    System.out.println("Pushing test set to Rhesis...");
    TestSet pushedTestSet = client.testSets().create(generatedTestSet);
    System.out.println("Successfully pushed TestSet! ID: " + pushedTestSet.id());

    // Create a dummy file to attach to each test
    Path tempFile = Files.createTempFile("company_policy", ".txt");
    Files.writeString(
        tempFile, "Company Policy: Employees are entitled to 20 vacation days per year.");
    System.out.println("Attaching file to each generated test...");
    // Retrieve the tests that were actually created on the backend
    List<Test> remoteTests = client.testSets().getTests(pushedTestSet.id());
    System.out.println("Fetched " + remoteTests.size() + " tests from the backend.");

    for (Test test : remoteTests) {
      if (test.id() != null) {
        System.out.println("Attaching file to test: " + test.id());
        client.tests().addFile(test.id(), tempFile);
      }
    }

    System.out.println("Successfully attached files!");

    // Clean up our local temporary file
    Files.deleteIfExists(tempFile);
  }
}