Add bandit support

build.gradle

@@ -6,7 +6,7 @@ plugins {
 }
 
 group = 'cloud.eppo'
-version = '2.1.0-SNAPSHOT'
+version = '3.0.0-SNAPSHOT'
 ext.isReleaseVersion = !version.endsWith("SNAPSHOT")
 
 dependencies {
@@ -23,6 +23,7 @@ dependencies {
   testImplementation 'commons-io:commons-io:2.11.0'
   testImplementation "com.google.truth:truth:1.4.4"
   testImplementation 'org.mockito:mockito-core:4.11.0'
+  testImplementation 'org.mockito:mockito-inline:4.11.0'
 }
 
 test {
@@ -140,14 +141,15 @@ tasks.withType(PublishToMavenRepository) {
   }
 }
 
-signing {
-  sign publishing.publications.mavenJava
-  if (System.env['CI']) {
-    useInMemoryPgpKeys(System.env.GPG_PRIVATE_KEY, System.env.GPG_PASSPHRASE)
+if (!project.gradle.startParameter.taskNames.contains('publishToMavenLocal')) {
+  signing {
+    sign publishing.publications.mavenJava
+    if (System.env['CI']) {
+      useInMemoryPgpKeys(System.env.GPG_PRIVATE_KEY, System.env.GPG_PASSPHRASE)
+    }
   }
 }
 
-
 javadoc {
   failOnError = false
   options.addStringOption('Xdoclint:none', '-quiet')

src/main/java/cloud/eppo/BanditEvaluationResult.java

@@ -0,0 +1,73 @@
+package cloud.eppo;
+
+import cloud.eppo.ufc.dto.DiscriminableAttributes;
+
+public class BanditEvaluationResult {
+
+  private final String flagKey;
+  private final String subjectKey;
+  private final DiscriminableAttributes subjectAttributes;
+  private final String actionKey;
+  private final DiscriminableAttributes actionAttributes;
+  private final double actionScore;
+  private final double actionWeight;
+  private final double gamma;
+  private final double optimalityGap;
+
+  public BanditEvaluationResult(
+      String flagKey,
+      String subjectKey,
+      DiscriminableAttributes subjectAttributes,
+      String actionKey,
+      DiscriminableAttributes actionAttributes,
+      double actionScore,
+      double actionWeight,
+      double gamma,
+      double optimalityGap) {
+    this.flagKey = flagKey;
+    this.subjectKey = subjectKey;
+    this.subjectAttributes = subjectAttributes;
+    this.actionKey = actionKey;
+    this.actionAttributes = actionAttributes;
+    this.actionScore = actionScore;
+    this.actionWeight = actionWeight;
+    this.gamma = gamma;
+    this.optimalityGap = optimalityGap;
+  }
+
+  public String getFlagKey() {
+    return flagKey;
+  }
+
+  public String getSubjectKey() {
+    return subjectKey;
+  }
+
+  public DiscriminableAttributes getSubjectAttributes() {
+    return subjectAttributes;
+  }
+
+  public String getActionKey() {
+    return actionKey;
+  }
+
+  public DiscriminableAttributes getActionAttributes() {
+    return actionAttributes;
+  }
+
+  public double getActionScore() {
+    return actionScore;
+  }
+
+  public double getActionWeight() {
+    return actionWeight;
+  }
+
+  public double getGamma() {
+    return gamma;
+  }
+
+  public double getOptimalityGap() {
+    return optimalityGap;
+  }
+}

src/main/java/cloud/eppo/BanditEvaluator.java

@@ -0,0 +1,168 @@
+package cloud.eppo;
+
+import cloud.eppo.ufc.dto.*;
+import java.util.*;
+import java.util.stream.Collectors;
+
+public class BanditEvaluator {
+
+  private static final int BANDIT_ASSIGNMENT_SHARDS = 10000; // hard-coded for now
+
+  public static BanditEvaluationResult evaluateBandit(
+      String flagKey,
+      String subjectKey,
+      DiscriminableAttributes subjectAttributes,
+      Actions actions,
+      BanditModelData modelData) {
+    Map<String, Double> actionScores = scoreActions(subjectAttributes, actions, modelData);
+    Map<String, Double> actionWeights =
+        weighActions(actionScores, modelData.getGamma(), modelData.getActionProbabilityFloor());
+    String selectedActionKey = selectAction(flagKey, subjectKey, actionWeights);
+
+    // Compute optimality gap in terms of score
+    double topScore =
+        actionScores.values().stream().mapToDouble(Double::doubleValue).max().orElse(0);
+    double optimalityGap = topScore - actionScores.get(selectedActionKey);
+
+    return new BanditEvaluationResult(
+        flagKey,
+        subjectKey,
+        subjectAttributes,
+        selectedActionKey,
+        actions.get(selectedActionKey),
+        actionScores.get(selectedActionKey),
+        actionWeights.get(selectedActionKey),
+        modelData.getGamma(),
+        optimalityGap);
+  }
+
+  private static Map<String, Double> scoreActions(
+      DiscriminableAttributes subjectAttributes, Actions actions, BanditModelData modelData) {
+    return actions.entrySet().stream()
+        .collect(
+            Collectors.toMap(
+                Map.Entry::getKey,
+                e -> {
+                  String actionName = e.getKey();
+                  DiscriminableAttributes actionAttributes = e.getValue();
+
+                  // get all coefficients known to the model for this action
+                  BanditCoefficients banditCoefficients =
+                      modelData.getCoefficients().get(actionName);
+
+                  if (banditCoefficients == null) {
+                    // Unknown action; return the default action score
+                    return modelData.getDefaultActionScore();
+                  }
+
+                  // Score the action using the provided attributes
+                  double actionScore = banditCoefficients.getIntercept();
+                  actionScore +=
+                      scoreContextForCoefficients(
+                          actionAttributes.getNumericAttributes(),
+                          banditCoefficients.getActionNumericCoefficients());
+                  actionScore +=
+                      scoreContextForCoefficients(
+                          actionAttributes.getCategoricalAttributes(),
+                          banditCoefficients.getActionCategoricalCoefficients());
+                  actionScore +=
+                      scoreContextForCoefficients(
+                          subjectAttributes.getNumericAttributes(),
+                          banditCoefficients.getSubjectNumericCoefficients());
+                  actionScore +=
+                      scoreContextForCoefficients(
+                          subjectAttributes.getCategoricalAttributes(),
+                          banditCoefficients.getSubjectCategoricalCoefficients());
+
+                  return actionScore;
+                }));
+  }
+
+  private static double scoreContextForCoefficients(
+      Attributes attributes, Map<String, ? extends BanditAttributeCoefficients> coefficients) {
+
+    double totalScore = 0.0;
+
+    for (BanditAttributeCoefficients attributeCoefficients : coefficients.values()) {
+      EppoValue contextValue = attributes.get(attributeCoefficients.getAttributeKey());
+      // The coefficient implementation knows how to score
+      double attributeScore = attributeCoefficients.scoreForAttributeValue(contextValue);
+      totalScore += attributeScore;
+    }
+
+    return totalScore;
+  }
+
+  private static Map<String, Double> weighActions(
+      Map<String, Double> actionScores, double gamma, double actionProbabilityFloor) {
+    Double highestScore = null;
+    String highestScoredAction = null;
+    for (Map.Entry<String, Double> actionScore : actionScores.entrySet()) {
+      if (highestScore == null
+          || actionScore.getValue() > highestScore
+          || actionScore
+                  .getValue()
+                  .equals(highestScore) // note: we break ties for scores by action name
+              && actionScore.getKey().compareTo(highestScoredAction) < 0) {
+        highestScore = actionScore.getValue();
+        highestScoredAction = actionScore.getKey();
+      }
+    }
+
+    // Weigh all the actions using their score
+    Map<String, Double> actionWeights = new HashMap<>();
+    double totalNonHighestWeight = 0.0;
+    for (Map.Entry<String, Double> actionScore : actionScores.entrySet()) {
+      if (actionScore.getKey().equals(highestScoredAction)) {
+        // The highest scored action is weighed at the end
+        continue;
+      }
+
+      // Compute weight (probability)
+      double unboundedProbability =
+          1 / (actionScores.size() + (gamma * (highestScore - actionScore.getValue())));
+      double minimumProbability = actionProbabilityFloor / actionScores.size();
+      double boundedProbability = Math.max(unboundedProbability, minimumProbability);
+      totalNonHighestWeight += boundedProbability;
+
+      actionWeights.put(actionScore.getKey(), boundedProbability);
+    }
+
+    // Weigh the highest scoring action (defensively preventing a negative probability)
+    double weightForHighestScore = Math.max(1 - totalNonHighestWeight, 0);
+    actionWeights.put(highestScoredAction, weightForHighestScore);
+    return actionWeights;
+  }
+
+  private static String selectAction(
+      String flagKey, String subjectKey, Map<String, Double> actionWeights) {
+    // Deterministically "shuffle" the actions
+    // This way as action weights shift, a bunch of users who were on the edge of one action won't
+    // all be shifted to the same new action at the same time.
+    List<String> shuffledActionKeys =
+        actionWeights.keySet().stream()
+            .sorted(
+                Comparator.comparingInt(
+                        (String actionKey) ->
+                            ShardUtils.getShard(
+                                flagKey + "-" + subjectKey + "-" + actionKey,
+                                BANDIT_ASSIGNMENT_SHARDS))
+                    .thenComparing(actionKey -> actionKey))
+            .collect(Collectors.toList());
+
+    // Select action from the shuffled actions, based on weight
+    double assignedShard =
+        ShardUtils.getShard(flagKey + "-" + subjectKey, BANDIT_ASSIGNMENT_SHARDS);
+    double assignmentWeightThreshold = assignedShard / (double) BANDIT_ASSIGNMENT_SHARDS;
+    double cumulativeWeight = 0;
+    String assignedAction = null;
+    for (String actionKey : shuffledActionKeys) {
+      cumulativeWeight += actionWeights.get(actionKey);
+      if (cumulativeWeight > assignmentWeightThreshold) {
+        assignedAction = actionKey;
+        break;
+      }
+    }
+    return assignedAction;
+  }
+}

src/main/java/cloud/eppo/ConfigurationRequestor.java

@@ -2,6 +2,8 @@
 
 import cloud.eppo.ufc.dto.FlagConfig;
 import java.io.IOException;
+import java.util.HashSet;
+import java.util.Set;
 import okhttp3.Response;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -12,28 +14,44 @@ public class ConfigurationRequestor {
 
   private final EppoHttpClient client;
   private final ConfigurationStore configurationStore;
+  private final Set<String> loadedBanditModelVersions;
 
   public ConfigurationRequestor(ConfigurationStore configurationStore, EppoHttpClient client) {
     this.configurationStore = configurationStore;
     this.client = client;
+    this.loadedBanditModelVersions = new HashSet<>();
   }
 
+  // TODO: async loading for android
   public void load() {
     log.debug("Fetching configuration");
-    Response response = client.get("/api/flag-config/v1/config");
+    String flagConfigurationJsonString = requestBody("/api/flag-config/v1/config");
+    configurationStore.setFlagsFromJsonString(flagConfigurationJsonString);
+
+    Set<String> neededModelVersions = configurationStore.banditModelVersions();
+    boolean needBanditParameters = !loadedBanditModelVersions.containsAll(neededModelVersions);
+    if (needBanditParameters) {
+      String banditParametersJsonString = requestBody("/api/flag-config/v1/bandits");
+      configurationStore.setBanditParametersFromJsonString(banditParametersJsonString);
+      // Record the model versions that we just loaded, so we can compare when the store is later
+      // updated
+      loadedBanditModelVersions.clear();
+      loadedBanditModelVersions.addAll(configurationStore.banditModelVersions());
+    }
+  }
+
+  private String requestBody(String route) {
+    Response response = client.get(route);
+    if (!response.isSuccessful() || response.body() == null) {
+      throw new RuntimeException("Failed to fetch from " + route);
+    }
     try {
-      if (!response.isSuccessful()) {
-        throw new RuntimeException("Failed to fetch configuration");
-      }
-      configurationStore.setFlagsFromJsonString(response.body().string());
+      return response.body().string();
     } catch (IOException e) {
-      // TODO: better exception handling?
       throw new RuntimeException(e);
     }
   }
 
-  // TODO: async loading for android
-
   public FlagConfig getConfiguration(String flagKey) {
     return configurationStore.getFlag(flagKey);
   }