store bandit flag to action association

src/client/eppo-client-with-bandits.spec.ts

@@ -13,14 +13,15 @@ import { IBanditEvent, IBanditLogger } from '../bandit-logger';
 import ConfigurationRequestor from '../configuration-requestor';
 import { MemoryOnlyConfigurationStore } from '../configuration-store/memory.store';
 import FetchHttpClient from '../http-client';
-import { BanditParameters, Flag } from '../interfaces';
+import { BanditFlagAssociation, BanditParameters, Flag } from '../interfaces';
 import { Attributes } from '../types';
 
 import EppoClient from './eppo-client';
 
 describe('EppoClient Bandits E2E test', () => {
   const flagStore = new MemoryOnlyConfigurationStore<Flag>();
-  const banditStore = new MemoryOnlyConfigurationStore<BanditParameters>();
+  const banditFlagStore = new MemoryOnlyConfigurationStore<BanditFlagAssociation[]>();
+  const banditModelStore = new MemoryOnlyConfigurationStore<BanditParameters>();
   let client: EppoClient;
   const mockLogAssignment = jest.fn();
   const mockLogBanditAction = jest.fn();
@@ -50,12 +51,17 @@ describe('EppoClient Bandits E2E test', () => {
       },
     });
     const httpClient = new FetchHttpClient(apiEndpoints, 1000);
-    const configurationRequestor = new ConfigurationRequestor(httpClient, flagStore, banditStore);
+    const configurationRequestor = new ConfigurationRequestor(
+      httpClient,
+      flagStore,
+      banditFlagStore,
+      banditModelStore,
+    );
     await configurationRequestor.fetchAndStoreConfigurations();
   });
 
   beforeEach(() => {
-    client = new EppoClient(flagStore, banditStore);
+    client = new EppoClient(flagStore, undefined, false, banditFlagStore, banditModelStore);
     client.setIsGracefulFailureMode(false);
     client.setAssignmentLogger({ logAssignment: mockLogAssignment });
     client.setBanditLogger({ logBanditAction: mockLogBanditAction });

src/client/eppo-client.spec.ts

@@ -33,7 +33,13 @@ export async function init(configurationStore: IConfigurationStore<Flag | Obfusc
     },
   });
   const httpClient = new FetchHttpClient(apiEndpoints, 1000);
-  const configurationRequestor = new ConfigurationRequestor(httpClient, configurationStore, null);
+  const configurationRequestor = new ConfigurationRequestor(
+    httpClient,
+    configurationStore,
+    // Leave bandit stores empty for this test
+    null,
+    null,
+  );
   await configurationRequestor.fetchAndStoreConfigurations();
 }
 
@@ -263,7 +269,7 @@ describe('EppoClient E2E test', () => {
     it.each(readTestData<IAssignmentTestCase>(ASSIGNMENT_TEST_DATA_DIR))(
       'test variation assignment splits',
       async ({ flag, variationType, defaultValue, subjects }: IAssignmentTestCase) => {
-        const client = new EppoClient(storage, undefined, undefined, true);
+        const client = new EppoClient(storage, undefined, true);
         client.setIsGracefulFailureMode(false);
 
         const typeAssignmentFunctions = {
@@ -616,7 +622,7 @@ describe('EppoClient E2E test', () => {
     });
 
     it('Fetches initial configuration with parameters in constructor', async () => {
-      client = new EppoClient(thisFlagStorage, undefined, requestConfiguration);
+      client = new EppoClient(thisFlagStorage, requestConfiguration);
       client.setIsGracefulFailureMode(false);
       // no configuration loaded
       let variation = client.getNumericAssignment(flagKey, subject, {}, 123.4);
@@ -647,7 +653,7 @@ describe('EppoClient E2E test', () => {
         }
       }
 
-      client = new EppoClient(new MockStore(), new MockStore(), requestConfiguration);
+      client = new EppoClient(new MockStore(), requestConfiguration);
       client.setIsGracefulFailureMode(false);
       // no configuration loaded
       let variation = client.getNumericAssignment(flagKey, subject, {}, 0.0);
@@ -689,7 +695,7 @@ describe('EppoClient E2E test', () => {
         ...requestConfiguration,
         pollAfterSuccessfulInitialization,
       };
-      client = new EppoClient(thisFlagStorage, undefined, requestConfiguration);
+      client = new EppoClient(thisFlagStorage, requestConfiguration);
       client.setIsGracefulFailureMode(false);
       // no configuration loaded
       let variation = client.getNumericAssignment(flagKey, subject, {}, 0.0);
@@ -754,7 +760,7 @@ describe('EppoClient E2E test', () => {
         throwOnFailedInitialization,
         pollAfterFailedInitialization,
       };
-      client = new EppoClient(thisFlagStorage, undefined, requestConfiguration);
+      client = new EppoClient(thisFlagStorage, requestConfiguration);
       client.setIsGracefulFailureMode(false);
       // no configuration loaded
       expect(client.getNumericAssignment(flagKey, subject, {}, 0.0)).toBe(0.0);

src/client/eppo-client.ts

@@ -21,7 +21,13 @@ import { decodeFlag } from '../decoding';
 import { EppoValue } from '../eppo_value';
 import { Evaluator, FlagEvaluation, noneResult } from '../evaluator';
 import FetchHttpClient from '../http-client';
-import { BanditParameters, Flag, ObfuscatedFlag, VariationType } from '../interfaces';
+import {
+  BanditFlagAssociation,
+  BanditParameters,
+  Flag,
+  ObfuscatedFlag,
+  VariationType,
+} from '../interfaces';
 import { getMD5Hash } from '../obfuscation';
 import initPoller, { IPoller } from '../poller';
 import { Attributes, AttributeType, ValueType } from '../types';
@@ -199,9 +205,11 @@ export default class EppoClient implements IEppoClient {
 
   constructor(
     private flagConfigurationStore: IConfigurationStore<Flag | ObfuscatedFlag>,
-    private banditConfigurationStore?: IConfigurationStore<BanditParameters>,
     private configurationRequestParameters?: FlagConfigurationRequestParameters,
     private isObfuscated = false,
+    // Parameter order is for backwards compatibility, and thus related items are not grouped
+    private banditFlagConfigurationStore?: IConfigurationStore<BanditFlagAssociation[]>,
+    private banditModelConfigurationStore?: IConfigurationStore<BanditParameters>,
   ) {}
 
   public setConfigurationRequestParameters(
@@ -216,10 +224,16 @@ export default class EppoClient implements IEppoClient {
     this.flagConfigurationStore = flagConfigurationStore;
   }
 
-  public setBanditConfigurationStore(
-    banditConfigurationStore: IConfigurationStore<BanditParameters>,
+  public setBanditFlagConfigurationStore(
+    banditFlagConfigurationStore: IConfigurationStore<BanditFlagAssociation[]>,
   ) {
-    this.banditConfigurationStore = banditConfigurationStore;
+    this.banditFlagConfigurationStore = banditFlagConfigurationStore;
+  }
+
+  public setBanditModelConfigurationStore(
+    banditModelConfigurationStore: IConfigurationStore<BanditParameters>,
+  ) {
+    this.banditModelConfigurationStore = banditModelConfigurationStore;
   }
 
   public setIsObfuscated(isObfuscated: boolean) {
@@ -267,7 +281,8 @@ export default class EppoClient implements IEppoClient {
     const configurationRequestor = new ConfigurationRequestor(
       httpClient,
       this.flagConfigurationStore,
-      this.banditConfigurationStore ?? null,
+      this.banditFlagConfigurationStore ?? null,
+      this.banditModelConfigurationStore ?? null,
     );
 
     this.requestPoller = initPoller(
@@ -393,14 +408,34 @@ export default class EppoClient implements IEppoClient {
     actions: Record<string, Attributes>, // TODO: ability to provide a set of actions with no context, or context broken out by numeric/categorical
     defaultValue: string,
   ): { variation: string; action: string | null } {
-    // TODO: store bandit flag information so that if no actions provided, we can quickly return with default value
-    let variation = this.getStringAssignment(flagKey, subjectKey, subjectAttributes, defaultValue);
+    let variation = defaultValue;
     let action: string | null = null;
-    const banditKey = variation;
     try {
-      const banditParameters = this.banditConfigurationStore?.get(banditKey);
-      if (banditParameters) {
-        // For now, we use the shortcut of assuming if a variation value is the key of a known bandit, that is the bandit we want
+      const banditFlagAssociations = this.banditFlagConfigurationStore?.get(flagKey);
+      if (banditFlagAssociations && !Object.keys(actions).length) {
+        // No actions passed for a flag known to have an active bandit, so we just return the default values so that
+        // we don't log a variation or bandit assignment
+        return { variation, action };
+      }
+
+      // Get the assigned variation for the flag with a possible bandit
+      variation = this.getStringAssignment(flagKey, subjectKey, subjectAttributes, defaultValue);
+
+      // Check if the assigned variation is an active bandit
+      // Note: the reason for non-bandit assignments include the subject being bucketed into a non-bandit variation or
+      // a rollout having been done.
+      const banditKey = banditFlagAssociations?.find(
+        (banditFlagAssociation) => banditFlagAssociation.variationValue === variation,
+      )?.key;
+
+      if (banditKey) {
+        // Retrieve the model parameters for the bandit
+        const banditParameters = this.banditModelConfigurationStore?.get(banditKey);
+
+        if (!banditParameters) {
+          throw new Error('No model parameters for bandit ' + banditKey);
+        }
+
         const banditModelData = banditParameters.modelData;
         const banditEvaluation = this.banditEvaluator.evaluateBandit(
           flagKey,
@@ -595,7 +630,8 @@ export default class EppoClient implements IEppoClient {
   public isInitialized() {
     return (
       this.flagConfigurationStore.isInitialized() &&
-      (!this.banditConfigurationStore || this.banditConfigurationStore.isInitialized())
+      (!this.banditFlagConfigurationStore || this.banditFlagConfigurationStore.isInitialized()) &&
+      (!this.banditModelConfigurationStore || this.banditModelConfigurationStore.isInitialized())
     );
   }
 

src/configuration-requestor.spec.ts

@@ -10,11 +10,12 @@ import ConfigurationRequestor from './configuration-requestor';
 import { IConfigurationStore } from './configuration-store/configuration-store';
 import { MemoryOnlyConfigurationStore } from './configuration-store/memory.store';
 import FetchHttpClient, { IHttpClient } from './http-client';
-import { BanditParameters, Flag } from './interfaces';
+import { BanditFlagAssociation, BanditParameters, Flag } from './interfaces';
 
 describe('ConfigurationRequestor', () => {
   let flagStore: IConfigurationStore<Flag>;
-  let banditStore: IConfigurationStore<BanditParameters>;
+  let banditFlagStore: IConfigurationStore<BanditFlagAssociation[]>;
+  let banditModelStore: IConfigurationStore<BanditParameters>;
   let httpClient: IHttpClient;
   let configurationRequestor: ConfigurationRequestor;
 
@@ -29,8 +30,14 @@ describe('ConfigurationRequestor', () => {
     });
     httpClient = new FetchHttpClient(apiEndpoints, 1000);
     flagStore = new MemoryOnlyConfigurationStore<Flag>();
-    banditStore = new MemoryOnlyConfigurationStore<BanditParameters>();
-    configurationRequestor = new ConfigurationRequestor(httpClient, flagStore, banditStore);
+    banditFlagStore = new MemoryOnlyConfigurationStore<BanditFlagAssociation[]>();
+    banditModelStore = new MemoryOnlyConfigurationStore<BanditParameters>();
+    configurationRequestor = new ConfigurationRequestor(
+      httpClient,
+      flagStore,
+      banditFlagStore,
+      banditModelStore,
+    );
   });
 
   afterEach(() => {
@@ -97,7 +104,7 @@ describe('ConfigurationRequestor', () => {
         end: 10000,
       });
 
-      expect(banditStore.getKeys().length).toBe(0);
+      expect(banditModelStore.getKeys().length).toBe(0);
     });
   });
 
@@ -129,9 +136,9 @@ describe('ConfigurationRequestor', () => {
       expect(flagStore.get('banner_bandit_flag')).toBeDefined();
       expect(flagStore.get('cold_start_bandit')).toBeDefined();
 
-      expect(banditStore.getKeys().length).toBeGreaterThanOrEqual(2);
+      expect(banditModelStore.getKeys().length).toBeGreaterThanOrEqual(2);
 
-      const bannerBandit = banditStore.get('banner_bandit');
+      const bannerBandit = banditModelStore.get('banner_bandit');
       expect(bannerBandit?.banditKey).toBe('banner_bandit');
       expect(bannerBandit?.modelName).toBe('falcon');
       expect(bannerBandit?.modelVersion).toBe('v123');
@@ -180,7 +187,7 @@ describe('ConfigurationRequestor', () => {
         bannerCoefficients['adidas'].subjectCategoricalCoefficients[0].valueCoefficients['female'],
       ).toBe(0);
 
-      const coldStartBandit = banditStore.get('cold_start_bandit');
+      const coldStartBandit = banditModelStore.get('cold_start_bandit');
       expect(coldStartBandit?.banditKey).toBe('cold_start_bandit');
       expect(coldStartBandit?.modelName).toBe('falcon');
       expect(coldStartBandit?.modelVersion).toBe('cold start');
@@ -192,7 +199,7 @@ describe('ConfigurationRequestor', () => {
     });
 
     it('Will not fetch bandit parameters if there is no store', async () => {
-      configurationRequestor = new ConfigurationRequestor(httpClient, flagStore, null);
+      configurationRequestor = new ConfigurationRequestor(httpClient, flagStore, null, null);
       await configurationRequestor.fetchAndStoreConfigurations();
       expect(fetchSpy).toHaveBeenCalledTimes(1);
     });

src/configuration-requestor.ts

@@ -1,13 +1,16 @@
 import { IConfigurationStore } from './configuration-store/configuration-store';
 import { IHttpClient } from './http-client';
-import { BanditParameters, Flag } from './interfaces';
+import { BanditFlagAssociation, BanditParameters, Flag } from './interfaces';
 
 // Requests AND stores flag configurations
 export default class ConfigurationRequestor {
   constructor(
     private readonly httpClient: IHttpClient,
     private readonly flagConfigurationStore: IConfigurationStore<Flag>,
-    private readonly banditConfigurationStore: IConfigurationStore<BanditParameters> | null,
+    private readonly flagBanditConfigurationStore: IConfigurationStore<
+      BanditFlagAssociation[]
+    > | null,
+    private readonly banditModelConfigurationStore: IConfigurationStore<BanditParameters> | null,
   ) {}
 
   async fetchAndStoreConfigurations(): Promise<void> {
@@ -18,16 +21,40 @@ export default class ConfigurationRequestor {
 
     await this.flagConfigurationStore.setEntries(configResponse.flags);
     const flagsHaveBandits = Object.keys(configResponse.bandits ?? {}).length > 0;
-    const banditStoreProvided = !!this.banditConfigurationStore;
-    if (flagsHaveBandits && banditStoreProvided) {
+    const banditStoresProvided = Boolean(
+      this.flagBanditConfigurationStore && this.banditModelConfigurationStore,
+    );
+    if (flagsHaveBandits && banditStoresProvided) {
+      // Map bandit flag associations by flag key for quick lookup (instead of bandit key as provided by the UFC)
+      const banditFlagAssociations = this.indexBanditFlagAssociationsByFlagKey(
+        configResponse.bandits,
+      );
+      this.flagBanditConfigurationStore?.setEntries(banditFlagAssociations);
       // TODO: different polling intervals for bandit parameters
       const banditResponse = await this.httpClient.getBanditParameters();
       if (banditResponse?.bandits) {
-        if (!this.banditConfigurationStore) {
+        if (!this.banditModelConfigurationStore) {
           throw new Error('Bandit parameters fetched but no bandit configuration store provided');
         }
-        await this.banditConfigurationStore.setEntries(banditResponse.bandits);
+        await this.banditModelConfigurationStore.setEntries(banditResponse.bandits);
       }
     }
   }
+
+  private indexBanditFlagAssociationsByFlagKey(
+    banditFlagAssociationsByBanditKey: Record<string, BanditFlagAssociation[]>,
+  ): Record<string, BanditFlagAssociation[]> {
+    const banditFlagAssociationsByFlagKey: Record<string, BanditFlagAssociation[]> = {};
+    Object.values(banditFlagAssociationsByBanditKey).forEach((banditFlags) => {
+      banditFlags.forEach((banditFlag) => {
+        let flagAssociations = banditFlagAssociationsByFlagKey[banditFlag.flagKey];
+        if (!flagAssociations) {
+          flagAssociations = [];
+          banditFlagAssociationsByFlagKey[banditFlag.flagKey] = flagAssociations;
+        }
+        flagAssociations.push(banditFlag);
+      });
+    });
+    return banditFlagAssociationsByFlagKey;
+  }
 }

src/http-client.ts

@@ -1,5 +1,5 @@
 import ApiEndpoints from './api-endpoints';
-import { BanditParameters, Flag } from './interfaces';
+import { BanditFlagAssociation, BanditParameters, Flag } from './interfaces';
 
 export interface IQueryParams {
   apiKey: string;
@@ -18,15 +18,7 @@ export class HttpRequestError extends Error {
 
 export interface IUniversalFlagConfigResponse {
   flags: Record<string, Flag>;
-  bandits: Record<
-    string,
-    {
-      key: string;
-      flagKey: string;
-      variationKey: string;
-      variationValue: string;
-    }
-  >;
+  bandits: Record<string, BanditFlagAssociation[]>;
 }
 
 export interface IBanditParametersResponse {

src/interfaces.ts

@@ -81,6 +81,13 @@ export interface ObfuscatedShard {
   ranges: Range[];
 }
 
+export interface BanditFlagAssociation {
+  key: string;
+  flagKey: string;
+  variationKey: string;
+  variationValue: string;
+}
+
 export interface BanditParameters {
   banditKey: string;
   modelName: string;