feat: typed Choice<E> and ChoiceAnswer<E> with enum constants as labels

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Typed choices: `Choice<E>` and `ChoiceAnswer<E>` carry the label type. `Choice.of(instructions, Dept.class)`, `Choice.builder(Dept.class)`, and `TypeSafeRequest.Builder.choice(key, Dept.class, …)` build a question from an enum's constants, and `response.choice(key, Dept.class)` reads the answer back with `choice()` as the enum and `probabilities()` keyed by it. String labels are unchanged: `Choice.of(…, String...)` infers `Choice<String>` and `response.choice(key)` returns `ChoiceAnswer<String>`. **Breaking:** `ChoiceAnswer` gained a type parameter, which changes the erased return type of `choice()`; code compiled against 0.4.0 or earlier must be recompiled, and a raw `ChoiceAnswer` needs `<String>` or `var` (#10).
 - A response missing its `usage` block, or a `usage` missing `input_tokens` or `output_tokens`, now fails `systemOne` with a `TypeSafeException` naming the field, instead of reading as `null` or `0` (#12).
 - `ScoreAnswer.legend()` is now `Map<String, Object>`, since levels are sent as any JSON value and echoed back as sent. A `Score` with an object or array level previously failed the whole response, including its other answers and usage, with `Cannot deserialize value of type java.lang.String` (#11).
 

README.md

@@ -76,7 +76,7 @@ TypeSafeResponse response = client.systemOne(r -> r
                 .level("Threatens loss or suspension within hours")));
 
 double phishing = response.noul("is_phishing");            // 0.0 to 1.0
-ChoiceAnswer category = response.choice("category");       // choice(), probabilities(), confidence()
+ChoiceAnswer<String> category = response.choice("category"); // choice(), probabilities(), confidence()
 ScoreAnswer urgency = response.score("urgency");           // score(), probabilities(), confidence(), legend()
 ```
 
@@ -91,12 +91,41 @@ at (`` `email.body` ``) rather than one long string. `state(key, value)` adds a
 ### Questions
 
 - `Noul.of(instructions)` asks yes or no; `whenTrue` and `whenFalse` describe the outcomes.
-- `Choice.of(instructions, labels...)` picks one label; `option(label, description)` describes a label, `option(label)` leaves it undescribed.
+- `Choice.of(instructions, labels...)` picks one label; `option(label, description)` describes a label, `option(label)` leaves it undescribed. Labels can also be the constants of an enum; see below.
 - `Score.of(instructions, levels...)` places the state on an ordered rubric of at least two levels.
 
 Instructions are optional when the criteria say enough on their own. Any description can be a plain string or a
 `Criterion` with `what`, `notFor`, and `examples`. Prebuilt questions are plain records and can be shared across requests.
 
+### Typed choices
+
+When the labels of a `Choice` are the constants of an enum you already have, build the question from the enum and read
+the answer back as that enum. A misspelled label is then a compile error, the probabilities are keyed by the constants,
+and a `switch` over the answer is exhaustive. The wire form is unchanged: the label is the constant's name.
+
+```java
+enum Dept { BILLING, SHIPPING, SECURITY }
+
+Choice<Dept> dept = Choice.of("Which team should handle `email`?", Dept.class);      // one option per constant
+Choice<Dept> described = Choice.builder(Dept.class)
+        .instructions("Which team should handle `email`?")
+        .option(Dept.BILLING, "Invoices, refunds, payment methods")
+        .option(Dept.SECURITY, o -> o.what("Credential theft").notFor("Legitimate requests"))
+        .build();                                                                     // only the constants named
+
+TypeSafeResponse response = client.systemOne(Map.of("email", email), Map.of("dept", dept));
+
+ChoiceAnswer<Dept> answer = response.choice("dept", Dept.class);
+answer.choice();                            // Dept.SECURITY
+answer.probabilities().get(Dept.BILLING);   // 0.48
+switch (answer.choice()) {                  // exhaustive: a missing case is a compile error
+    case BILLING -> ...; case SHIPPING -> ...; case SECURITY -> ...;
+}
+```
+
+`response.choice("dept")` still returns the `String` form. Reading an answer as an enum that lacks one of its labels throws
+an `IllegalArgumentException` naming the label and the enum's constants.
+
 ### Criteria-driven questions
 
 When the rules are user-defined data rather than code, `CriteriaQuestionSet` puts them into the state and generates one

typesafe-sdk/src/main/java/io/github/premocloud/typesafe/Choice.java

@@ -7,27 +7,38 @@
 import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.function.Consumer;
+import java.util.function.Function;
 
 /**
  * Pick one option from a named set. The answer carries a probability per option.
  *
+ * <p>Labels are plain strings, or the constants of an enum. An enum choice reads back typed through
+ * {@link TypeSafeResponse#choice(String, Class)}, so a misspelled label is a compile error and a switch over the
+ * answer is exhaustive. The wire form is the same either way: the label is the constant's name.
+ *
  * <pre>{@code
  * Choice.of("What is this ticket about?", "billing", "technical", "other")       // undescribed labels
+ * Choice.of("Which team should handle this?", Dept.class)                         // one option per constant
+ * Choice.builder(Dept.class).instructions("Which team should handle this?")
+ *         .option(Dept.BILLING, "Invoices, refunds, payment methods")
+ *         .option(Dept.SECURITY, o -> o.what("Credential theft").notFor("Legitimate requests"))
+ *         .build()
  * Choice.of(c -> c.instructions("Which category?")
  *         .option("MARKETING", "Promotional content sent to a list")
  *         .option("PHISHING", o -> o.what("Credential theft").notFor("Legitimate requests"))
  *         .option("OTHER"))
  * }</pre>
  *
+ * @param <E>          the label type: {@code String}, or an enum whose constant names are the labels
  * @param instructions the question, or {@code null}
  * @param criteria     labels mapped to a description, a {@link Criterion}, or {@code null} for an undescribed label
  */
 @JsonInclude(JsonInclude.Include.NON_NULL)
-public record Choice(@Nullable Object instructions, Map<String, @Nullable Object> criteria) implements TypeSafeQuestion {
+public record Choice<E>(@Nullable Object instructions, Map<String, @Nullable Object> criteria) implements TypeSafeQuestion {
 
     /** Labels without descriptions, as in {@code choice("Which?", {billing: null, technical: null})}. */
-    public static Choice of(Object instructions, String... labels) {
-        Builder builder = builder().instructions(instructions);
+    public static Choice<String> of(Object instructions, String... labels) {
+        Builder<String> builder = builder().instructions(instructions);
 
         for (String label : labels) {
             builder.option(label);
@@ -36,51 +47,72 @@ public static Choice of(Object instructions, String... labels) {
         return builder.build();
     }
 
-    public static Choice of(Object instructions, Map<String, ?> criteria) {
-        Builder builder = builder().instructions(instructions);
+    public static Choice<String> of(Object instructions, Map<String, ?> criteria) {
+        Builder<String> builder = builder().instructions(instructions);
         criteria.forEach(builder::option);
         return builder.build();
     }
 
-    public static Choice of(Consumer<Builder> configure) {
-        Builder builder = builder();
+    /** One undescribed option per constant of {@code labels}, in declaration order. */
+    public static <E extends Enum<E>> Choice<E> of(Object instructions, Class<E> labels) {
+        Builder<E> builder = builder(labels).instructions(instructions);
+
+        for (E constant : labels.getEnumConstants()) {
+            builder.option(constant);
+        }
+
+        return builder.build();
+    }
+
+    public static Choice<String> of(Consumer<Builder<String>> configure) {
+        Builder<String> builder = builder();
         configure.accept(builder);
         return builder.build();
     }
 
-    public static Builder builder() {
-        return new Builder();
+    public static Builder<String> builder() {
+        return new Builder<>(Function.identity());
+    }
+
+    /** A builder whose options are constants of {@code labels}; add the ones to ask about with {@code option}. */
+    public static <E extends Enum<E>> Builder<E> builder(Class<E> labels) {
+        return new Builder<>(Enum::name);
     }
 
-    public static final class Builder {
+    public static final class Builder<E> {
+        private final Function<E, String> label;
         private @Nullable Object instructions;
         private final Map<String, @Nullable Object> options = new LinkedHashMap<>();
 
-        public Builder instructions(Object instructions) {
+        private Builder(Function<E, String> label) {
+            this.label = label;
+        }
+
+        public Builder<E> instructions(Object instructions) {
             this.instructions = instructions;
             return this;
         }
 
         /** An undescribed label. */
-        public Builder option(String label) {
+        public Builder<E> option(E label) {
             return option(label, (Object) null);
         }
 
-        public Builder option(String label, @Nullable Object description) {
-            options.put(label, description);
+        public Builder<E> option(E label, @Nullable Object description) {
+            options.put(this.label.apply(label), description);
             return this;
         }
 
-        public Builder option(String label, Consumer<Criterion.Builder> configure) {
+        public Builder<E> option(E label, Consumer<Criterion.Builder> configure) {
             return option(label, Criterion.of(configure));
         }
 
-        public Choice build() {
+        public Choice<E> build() {
             if (options.isEmpty()) {
                 throw new IllegalStateException("A choice question needs at least one option");
             }
 
-            return new Choice(instructions, Collections.unmodifiableMap(new LinkedHashMap<>(options)));
+            return new Choice<>(instructions, Collections.unmodifiableMap(new LinkedHashMap<>(options)));
         }
     }
 }

typesafe-sdk/src/main/java/io/github/premocloud/typesafe/ChoiceAnswer.java

@@ -4,26 +4,56 @@
 import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
 import com.fasterxml.jackson.annotation.JsonProperty;
 
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.EnumMap;
 import java.util.Map;
 
 /**
+ * @param <E>           the label type: {@code String} as read from the response, or an enum after {@link #as(Class)}
  * @param choice        highest-probability option
  * @param probabilities probability per option, summing to 1
  * @param confidence    0 to 1, how concentrated the distribution is
  */
 @JsonIgnoreProperties(ignoreUnknown = true)
-public record ChoiceAnswer(String choice, Map<String, Double> probabilities, double confidence) implements TypeSafeAnswer {
+public record ChoiceAnswer<E>(E choice, Map<E, Double> probabilities, double confidence) implements TypeSafeAnswer {
 
-    /** Jackson entry point: a choice answer missing any of its fields is malformed. */
+    /** Jackson entry point: a choice answer missing any of its fields is malformed. Labels come off the wire as Strings. */
     @JsonCreator
+    @SuppressWarnings("unchecked")
     ChoiceAnswer(
             @JsonProperty("choice") String choice,
             @JsonProperty("probabilities") Map<String, Double> probabilities,
             @JsonProperty("confidence") Double confidence,
             @JsonProperty("type") String ignoredType
     ) {
-        this(TypeSafeAnswer.required(choice, "choice", "choice"),
-                TypeSafeAnswer.required(probabilities, "choice", "probabilities"),
+        this((E) TypeSafeAnswer.required(choice, "choice", "choice"),
+                (Map<E, Double>) TypeSafeAnswer.required(probabilities, "choice", "probabilities"),
                 TypeSafeAnswer.required(confidence, "choice", "confidence").doubleValue());
     }
+
+    /**
+     * The same answer with its labels as constants of {@code labels}, matched by name. The probabilities come back in
+     * the enum's declaration order.
+     *
+     * @throws IllegalArgumentException if the chosen label or any probability key is not a constant of {@code labels}
+     */
+    public <T extends Enum<T>> ChoiceAnswer<T> as(Class<T> labels) {
+        Map<T, Double> typed = new EnumMap<>(labels);
+        probabilities.forEach((label, probability) -> typed.put(constant(labels, label), probability));
+        return new ChoiceAnswer<>(constant(labels, choice), Collections.unmodifiableMap(typed), confidence);
+    }
+
+    private static <T extends Enum<T>> T constant(Class<T> labels, Object label) {
+        String name = String.valueOf(label);
+
+        for (T constant : labels.getEnumConstants()) {
+            if (constant.name().equals(name)) {
+                return constant;
+            }
+        }
+
+        throw new IllegalArgumentException("Choice '%s' is not a constant of %s; expected one of %s"
+                .formatted(name, labels.getSimpleName(), Arrays.toString(labels.getEnumConstants())));
+    }
 }

typesafe-sdk/src/main/java/io/github/premocloud/typesafe/TypeSafeRequest.java

@@ -79,10 +79,17 @@ public Builder noul(String key, Consumer<Noul.Builder> configure) {
             return question(key, Noul.of(configure));
         }
 
-        public Builder choice(String key, Consumer<Choice.Builder> configure) {
+        public Builder choice(String key, Consumer<Choice.Builder<String>> configure) {
             return question(key, Choice.of(configure));
         }
 
+        /** A choice whose labels are constants of {@code labels}; add the ones to ask about with {@code option}. */
+        public <E extends Enum<E>> Builder choice(String key, Class<E> labels, Consumer<Choice.Builder<E>> configure) {
+            Choice.Builder<E> builder = Choice.builder(labels);
+            configure.accept(builder);
+            return question(key, builder.build());
+        }
+
         public Builder score(String key, Consumer<Score.Builder> configure) {
             return question(key, Score.of(configure));
         }

typesafe-sdk/src/main/java/io/github/premocloud/typesafe/TypeSafeResponse.java

@@ -26,10 +26,20 @@ public double noul(String key) {
         return answer(key, NoulAnswer.class).noul();
     }
 
-    public ChoiceAnswer choice(String key) {
+    /** @return the choice answer with its labels as Strings, as they came off the wire */
+    @SuppressWarnings("unchecked")
+    public ChoiceAnswer<String> choice(String key) {
         return answer(key, ChoiceAnswer.class);
     }
 
+    /**
+     * @return the choice answer with its labels as constants of {@code labels}, for a question built from that enum
+     * @throws IllegalArgumentException if a label in the answer is not a constant of {@code labels}
+     */
+    public <E extends Enum<E>> ChoiceAnswer<E> choice(String key, Class<E> labels) {
+        return choice(key).as(labels);
+    }
+
     public ScoreAnswer score(String key) {
         return answer(key, ScoreAnswer.class);
     }
@@ -38,8 +48,9 @@ public Map<String, NoulAnswer> nouls() {
         return answersOf(NoulAnswer.class);
     }
 
-    public Map<String, ChoiceAnswer> choices() {
-        return answersOf(ChoiceAnswer.class);
+    @SuppressWarnings({"unchecked", "rawtypes"})
+    public Map<String, ChoiceAnswer<String>> choices() {
+        return (Map) answersOf(ChoiceAnswer.class);
     }
 
     public Map<String, ScoreAnswer> scores() {

typesafe-sdk/src/test/java/io/github/premocloud/typesafe/TypeSafeClientTest.java

@@ -85,7 +85,7 @@ void systemOnePostsBearerAuthenticatedJsonAndReturnsTypedAnswers() throws Except
 
         assertEquals("jev-1.13.0", response.model());
         assertEquals(0.93, response.noul("is_phishing"));
-        ChoiceAnswer category = response.choice("spam_category");
+        ChoiceAnswer<String> category = response.choice("spam_category");
         assertEquals("PHISHING", category.choice());
         assertEquals(0.9, category.probabilities().get("PHISHING"));
         assertEquals(0.88, category.confidence());
@@ -310,6 +310,39 @@ void systemOneRejectsNoulAnswerWithExplicitNullValue() {
         assertTrue(exception.getMessage().contains("noul"), exception.getMessage());
     }
 
+    /** Declared in the opposite order to the response's probabilities, to show the typed map follows the enum. */
+    enum SpamCategory { MARKETING, PHISHING }
+
+    enum Urgency { LOW, HIGH }
+
+    @Test
+    void systemOneReadsAChoiceAnswerAsAnEnum() {
+        server.reply(200, RESPONSE_JSON);
+
+        TypeSafeResponse response = client.systemOne(spamRequest());
+
+        ChoiceAnswer<SpamCategory> category = response.choice("spam_category", SpamCategory.class);
+        assertEquals(SpamCategory.PHISHING, category.choice());
+        assertEquals(0.9, category.probabilities().get(SpamCategory.PHISHING));
+        assertEquals(0.1, category.probabilities().get(SpamCategory.MARKETING));
+        assertEquals(0.88, category.confidence());
+        assertEquals(List.of(SpamCategory.MARKETING, SpamCategory.PHISHING), List.copyOf(category.probabilities().keySet()));
+        assertEquals("PHISHING", response.choice("spam_category").choice(), "the String read is unchanged");
+        assertEquals("PHISHING", response.choices().get("spam_category").choice());
+    }
+
+    @Test
+    void systemOneRejectsAChoiceLabelThatIsNotAConstantOfTheEnum() {
+        server.reply(200, RESPONSE_JSON);
+        TypeSafeResponse response = client.systemOne(spamRequest());
+
+        IllegalArgumentException exception = assertThrows(IllegalArgumentException.class,
+                () -> response.choice("spam_category", Urgency.class));
+
+        assertTrue(exception.getMessage().contains("PHISHING"), exception.getMessage());
+        assertTrue(exception.getMessage().contains("[LOW, HIGH]"), exception.getMessage());
+    }
+
     @Test
     void systemOneRejectsChoiceAnswerMissingItsChoice() {
         server.reply(200, """

typesafe-sdk/src/test/java/io/github/premocloud/typesafe/TypeSafeRequestTest.java

@@ -6,6 +6,7 @@
 
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
@@ -127,6 +128,35 @@ void structuredStateAndInstructionsPassThrough() {
         assertThrows(IllegalStateException.class, () -> TypeSafeRequest.builder().state("text").state("k", "v"));
     }
 
+    enum Dept { BILLING, SHIPPING, SECURITY }
+
+    @Test
+    void enumChoicesUseTheConstantNamesAsLabelsAndKeepTheWireShape() throws Exception {
+        Choice<Dept> flat = Choice.of("Which team?", Dept.class);
+        Choice<Dept> described = Choice.builder(Dept.class).instructions("Which team?")
+                .option(Dept.BILLING, "Invoices and refunds").option(Dept.SECURITY).build();
+        TypeSafeRequest request = TypeSafeRequest.of(r -> r.state("text")
+                .choice("dept", Dept.class, c -> c.instructions("Which team?").option(Dept.SHIPPING, o -> o.what("Delivery"))));
+
+        assertEquals(List.of("BILLING", "SHIPPING", "SECURITY"), List.copyOf(flat.criteria().keySet()));
+        JsonNode flatJson = objectMapper.valueToTree(flat);
+        assertEquals(Set.of("type", "instructions", "criteria"), fieldNames(flatJson), "no enum metadata leaks onto the wire");
+        assertEquals(objectMapper.valueToTree(Choice.of("Which team?", "BILLING", "SHIPPING", "SECURITY")), flatJson);
+
+        JsonNode describedJson = objectMapper.valueToTree(described);
+        assertEquals("Invoices and refunds", describedJson.at("/criteria/BILLING").asText());
+        assertTrue(describedJson.at("/criteria/SECURITY").isNull());
+        assertFalse(describedJson.at("/criteria").has("SHIPPING"));
+
+        assertEquals("Delivery", objectMapper.valueToTree(request).at("/questions/dept/criteria/SHIPPING/what").asText());
+    }
+
+    private static Set<String> fieldNames(JsonNode node) {
+        Set<String> names = new java.util.HashSet<>();
+        node.fieldNames().forEachRemaining(names::add);
+        return names;
+    }
+
     @Test
     void prebuiltQuestionsCanBeReused() {
         Noul shared = Noul.of("Is it urgent?");