Log experiments with REST API

1
public static void main(String[] args) {
2
3
    // This example uses Jackson to handle JSON serialization and deserialization.
4
    ObjectMapper mapper = new ObjectMapper();
5
6
    String baseURI = System.getenv("OPIK_URL_OVERRIDE"); // e.g., "http://localhost:5173/api"
7
    String workspaceName = System.getenv("WORKSPACE_NAME");
8
    String apiKey = System.getenv("API_KEY");
9
10
    String datasetName = "my-dataset";
11
    String experimentName = "my-experiment";
12
13
    Map<String, String> requestBody = Map.of("dataset_name", datasetName);
14
15
    try (var client = HttpClient.newHttpClient()) {
16
17
        var request = HttpRequest.newBuilder()
18
                .uri(URI.create(baseURI).resolve("/v1/private/datasets/items/stream"))
19
                .header("Content-Type", "application/json")
20
                .header("Accept", "application/octet-stream")
21
                .header("Authorization", apiKey)
22
                .header("Comet-Workspace", workspaceName)
23
                .POST(HttpRequest.BodyPublishers.ofString(mapper.writeValueAsString(requestBody)))
24
                .build();
25
26
        HttpResponse<InputStream> response = client.send(
27
                request,
28
                HttpResponse.BodyHandlers.ofInputStream()
29
        );
30
31
        List<JsonNode> experimentItems = new ArrayList<>();
32
        try (var reader = new BufferedReader(new InputStreamReader(response.body()))) {
33
            String line;
34
            while ((line = reader.readLine()) != null) {
35
                JsonNode datasetItem = mapper.readTree(line);
36
                String question = datasetItem.get("data").get("question").asText();
37
                UUID datasetItemId = UUID.fromString(datasetItem.get("id").asText());
38
39
                JsonNode evaluationTaskOutput = callTargetApplication(question);
40
                List<JsonNode> scoresResults = callLmmOrMetricCalc(evaluationTaskOutput);
41
42
                ArrayNode scoreResultsArray = JsonNodeFactory.instance.arrayNode().addAll(scoresResults);
43
44
                JsonNode experimentItem = JsonNodeFactory.instance.objectNode()
45
                        .put("dataset_item_id", datasetItemId.toString())
46
                        .setAll(Map.of(
47
                                "evaluation_task_output", evaluationTaskOutput,
48
                                "feedback_scores", scoreResultsArray));
49
50
                experimentItems.add(experimentItem);
51
            }
52
        }
53
54
        var experimentBatchBody = JsonNodeFactory.instance.objectNode()
55
                .put("dataset_name", datasetName)
56
                .put("experiment_name", experimentName)
57
                .setAll(Map.of("items", JsonNodeFactory.instance.arrayNode().addAll(experimentItems)));
58
59
        var experimentRequest = HttpRequest.newBuilder()
60
                .uri(URI.create(baseURI).resolve("/v1/private/experiments/items/bulk"))
61
                .header("Content-Type", "application/json")
62
                .header("Authorization", apiKey)
63
                .header("Comet-Workspace", workspaceName)
64
                .PUT(HttpRequest.BodyPublishers.ofString(experimentBatchBody.toString()))
65
                .build();
66
67
        HttpResponse<String> experimentResponse = client.send(experimentRequest, HttpResponse.BodyHandlers.ofString());
68
69
        if (experimentResponse.statusCode() == 204) {
70
            System.out.println("Experiment items successfully created.");
71
        } else {
72
            System.out.printf("Failed to create experiment items, status code %s body %s",
73
                experimentResponse.statusCode(), experimentResponse.body());
74
        }
75
76
    } catch (InterruptedException | IOException e) {
77
        throw new RuntimeException(e);
78
    }
79
}