feat: evaluation v2 (#36)

amplitude · Dec 1, 2023 · 5443a7a · 5443a7a
1 parent 3bbada2
commit 5443a7a
Show file tree

Hide file tree

Showing 21 changed files with 748 additions and 363 deletions.
diff --git a/packages/node/package.json b/packages/node/package.json
@@ -32,6 +32,6 @@
   "dependencies": {
     "@amplitude/analytics-node": "^1.3.4",
     "@amplitude/analytics-types": "^1.3.1",
-    "@amplitude/evaluation-js": "1.1.1"
+    "@amplitude/experiment-core": "^0.7.1"
   }
 }
diff --git a/packages/node/src/assignment/assignment-service.ts b/packages/node/src/assignment/assignment-service.ts
@@ -7,7 +7,6 @@ import { Assignment, AssignmentFilter, AssignmentService } from './assignment';
 
 export const DAY_MILLIS = 24 * 60 * 60 * 1000;
 export const FLAG_TYPE_MUTUAL_EXCLUSION_GROUP = 'mutual-exclusion-group';
-export const FLAG_TYPE_HOLDOUT_GROUP = 'holdout-group';
 
 export class AmplitudeAssignmentService implements AssignmentService {
   private readonly amplitude: CoreClient;
@@ -20,44 +19,48 @@ export class AmplitudeAssignmentService implements AssignmentService {
 
   async track(assignment: Assignment): Promise<void> {
     if (this.assignmentFilter.shouldTrack(assignment)) {
-      this.amplitude.logEvent(this.toEvent(assignment));
+      this.amplitude.logEvent(toEvent(assignment));
     }
   }
+}
 
-  public toEvent(assignment: Assignment): BaseEvent {
-    const event: BaseEvent = {
-      event_type: '[Experiment] Assignment',
-      user_id: assignment.user.user_id,
-      device_id: assignment.user.device_id,
-      event_properties: {},
-      user_properties: {},
-    };
-
-    for (const resultsKey in assignment.results) {
-      event.event_properties[`${resultsKey}.variant`] =
-        assignment.results[resultsKey].value;
+export const toEvent = (assignment: Assignment): BaseEvent => {
+  const event: BaseEvent = {
+    event_type: '[Experiment] Assignment',
+    user_id: assignment.user.user_id,
+    device_id: assignment.user.device_id,
+    event_properties: {},
+    user_properties: {},
+  };
+  const set = {};
+  const unset = {};
+  for (const flagKey in assignment.results) {
+    const variant = assignment.results[flagKey];
+    if (!variant.key) {
+      continue;
     }
-
-    const set = {};
-    const unset = {};
-    for (const resultsKey in assignment.results) {
-      if (
-        assignment.results[resultsKey].type == FLAG_TYPE_MUTUAL_EXCLUSION_GROUP
-      ) {
-        continue;
-      } else if (assignment.results[resultsKey].isDefaultVariant) {
-        unset[`[Experiment] ${resultsKey}`] = '-';
+    const version = variant.metadata?.flagVersion;
+    const segmentName = variant.metadata?.segmentName;
+    const flagType = variant.metadata?.flagType;
+    const isDefault: boolean = variant.metadata?.default as boolean;
+    event.event_properties[`${flagKey}.variant`] = variant.key;
+    if (version && segmentName) {
+      event.event_properties[
+        `${flagKey}.details`
+      ] = `v${version} rule:${segmentName}`;
+    }
+    if (flagType != FLAG_TYPE_MUTUAL_EXCLUSION_GROUP) {
+      if (isDefault) {
+        unset[`[Experiment] ${flagKey}`] = '-';
       } else {
-        set[`[Experiment] ${resultsKey}`] =
-          assignment.results[resultsKey].value;
+        set[`[Experiment] ${flagKey}`] = variant.key;
       }
     }
-    event.user_properties['$set'] = set;
-    event.user_properties['$unset'] = unset;
-
-    event.insert_id = `${event.user_id} ${event.device_id} ${hashCode(
-      assignment.canonicalize(),
-    )} ${Math.floor(assignment.timestamp / DAY_MILLIS)}`;
-    return event;
   }
-}
+  event.user_properties['$set'] = set;
+  event.user_properties['$unset'] = unset;
+  event.insert_id = `${event.user_id} ${event.device_id} ${hashCode(
+    assignment.canonicalize(),
+  )} ${Math.floor(assignment.timestamp / DAY_MILLIS)}`;
+  return event;
+};
diff --git a/packages/node/src/assignment/assignment.ts b/packages/node/src/assignment/assignment.ts
@@ -1,5 +1,6 @@
+import { EvaluationVariant } from '@amplitude/experiment-core';
+
 import { ExperimentUser } from '../types/user';
-import { Results } from '../types/variant';
 
 export interface AssignmentService {
   track(assignment: Assignment): Promise<void>;
@@ -11,19 +12,24 @@ export interface AssignmentFilter {
 
 export class Assignment {
   public user: ExperimentUser;
-  public results: Results;
+  public results: Record<string, EvaluationVariant>;
   public timestamp: number = Date.now();
 
-  public constructor(user: ExperimentUser, results: Results) {
+  public constructor(
+    user: ExperimentUser,
+    results: Record<string, EvaluationVariant>,
+  ) {
     this.user = user;
     this.results = results;
   }
 
   public canonicalize(): string {
     let canonical = `${this.user.user_id?.trim()} ${this.user.device_id?.trim()} `;
     for (const key of Object.keys(this.results).sort()) {
-      const value = this.results[key];
-      canonical += key.trim() + ' ' + value?.value?.trim() + ' ';
+      const variant = this.results[key];
+      if (variant?.key) {
+        canonical += key.trim() + ' ' + variant?.key?.trim() + ' ';
+      }
     }
     return canonical;
   }

diff --git a/packages/node/src/local/cache.ts b/packages/node/src/local/cache.ts
@@ -1,11 +1,21 @@
 import { FlagConfigCache, FlagConfig } from '../types/flag';
 
 export class InMemoryFlagConfigCache implements FlagConfigCache {
-  private cache: Record<string, FlagConfig> = {};
+  private readonly store: FlagConfigCache | undefined;
+  private cache: Record<string, FlagConfig>;
 
-  public constructor(flagConfigs: Record<string, FlagConfig> = {}) {
+  public constructor(
+    store?: FlagConfigCache,
+    flagConfigs: Record<string, FlagConfig> = {},
+  ) {
+    this.store = store;
     this.cache = flagConfigs;
   }
+
+  public getAllCached(): Record<string, FlagConfig> {
+    return { ...this.cache };
+  }
+
   public async get(flagKey: string): Promise<FlagConfig> {
     return this.cache[flagKey];
   }
@@ -14,6 +24,7 @@ export class InMemoryFlagConfigCache implements FlagConfigCache {
   }
   public async put(flagKey: string, flagConfig: FlagConfig): Promise<void> {
     this.cache[flagKey] = flagConfig;
+    await this.store?.put(flagKey, flagConfig);
   }
   public async putAll(flagConfigs: Record<string, FlagConfig>): Promise<void> {
     for (const key in flagConfigs) {
@@ -22,11 +33,14 @@ export class InMemoryFlagConfigCache implements FlagConfigCache {
         this.cache[key] = flag;
       }
     }
+    await this.store?.putAll(flagConfigs);
   }
   public async delete(flagKey: string): Promise<void> {
     delete this.cache[flagKey];
+    await this.store?.delete(flagKey);
   }
   public async clear(): Promise<void> {
     this.cache = {};
+    await this.store?.clear();
   }
 }
diff --git a/packages/node/src/local/client.ts b/packages/node/src/local/client.ts
@@ -1,26 +1,31 @@
 import * as amplitude from '@amplitude/analytics-node';
-import evaluation from '@amplitude/evaluation-js';
+import {
+  EvaluationEngine,
+  EvaluationFlag,
+  topologicalSort,
+} from '@amplitude/experiment-core';
 
 import { Assignment, AssignmentService } from '../assignment/assignment';
 import { InMemoryAssignmentFilter } from '../assignment/assignment-filter';
-import {
-  AmplitudeAssignmentService,
-  FLAG_TYPE_HOLDOUT_GROUP,
-  FLAG_TYPE_MUTUAL_EXCLUSION_GROUP,
-} from '../assignment/assignment-service';
+import { AmplitudeAssignmentService } from '../assignment/assignment-service';
 import { FetchHttpClient } from '../transport/http';
 import {
   AssignmentConfig,
   AssignmentConfigDefaults,
   LocalEvaluationConfig,
   LocalEvaluationDefaults,
 } from '../types/config';
-import { FlagConfig, FlagConfigCache } from '../types/flag';
+import { FlagConfigCache } from '../types/flag';
 import { HttpClient } from '../types/transport';
 import { ExperimentUser } from '../types/user';
-import { Results, Variants } from '../types/variant';
+import { Variant, Variants } from '../types/variant';
 import { ConsoleLogger } from '../util/logger';
 import { Logger } from '../util/logger';
+import { convertUserToEvaluationContext } from '../util/user';
+import {
+  evaluationVariantsToVariants,
+  filterDefaultVariants,
+} from '../util/variant';
 
 import { InMemoryFlagConfigCache } from './cache';
 import { FlagConfigFetcher } from './fetcher';
@@ -34,22 +39,20 @@ export class LocalEvaluationClient {
   private readonly logger: Logger;
   private readonly config: LocalEvaluationConfig;
   private readonly poller: FlagConfigPoller;
-  private flags: FlagConfig[];
   private readonly assignmentService: AssignmentService;
+  private readonly evaluation: EvaluationEngine;
 
   /**
    * Directly access the client's flag config cache.
    *
    * Used for directly manipulating the flag configs used for evaluation.
    */
-  public readonly cache: FlagConfigCache;
+  public readonly cache: InMemoryFlagConfigCache;
 
   constructor(
     apiKey: string,
     config: LocalEvaluationConfig,
-    flagConfigCache: FlagConfigCache = new InMemoryFlagConfigCache(
-      config?.bootstrap,
-    ),
+    flagConfigCache?: FlagConfigCache,
     httpClient: HttpClient = new FetchHttpClient(config?.httpAgent),
   ) {
     this.config = { ...LocalEvaluationDefaults, ...config };
@@ -59,11 +62,10 @@ export class LocalEvaluationClient {
       this.config.serverUrl,
       this.config.debug,
     );
-    // We no longer use the flag config cache for accessing variants.
-    fetcher.setRawReceiver((flags: string) => {
-      this.flags = JSON.parse(flags);
-    });
-    this.cache = flagConfigCache;
+    this.cache = new InMemoryFlagConfigCache(
+      flagConfigCache,
+      this.config.bootstrap,
+    );
     this.logger = new ConsoleLogger(this.config.debug);
     this.poller = new FlagConfigPoller(
       fetcher,
@@ -80,6 +82,7 @@ export class LocalEvaluationClient {
         this.config.assignmentConfig,
       );
     }
+    this.evaluation = new EvaluationEngine();
   }
 
   private createAssignmentService(
@@ -94,6 +97,36 @@ export class LocalEvaluationClient {
     );
   }
 
+  /**
+   * Locally evaluate varints for a user.
+   *
+   * This function will only evaluate flags for the keys specified in the
+   * {@link flagKeys} argument. If {@link flagKeys} is missing, all flags in the
+   * {@link FlagConfigCache} will be evaluated.
+   *
+   * Unlike {@link evaluate}, this function returns a default variant object
+   * if the flag or experiment was evaluated, but the user was not assigned a
+   * variant (i.e. 'off').
+   *
+   * @param user The user to evaluate
+   * @param flagKeys The flags to evaluate with the user. If empty, all flags
+   * from the flag cache are evaluated.
+   * @returns The evaluated variants
+   */
+  public evaluateV2(
+    user: ExperimentUser,
+    flagKeys?: string[],
+  ): Record<string, Variant> {
+    const flags = this.cache.getAllCached() as Record<string, EvaluationFlag>;
+    this.logger.debug('[Experiment] evaluate - user:', user, 'flags:', flags);
+    const context = convertUserToEvaluationContext(user);
+    const sortedFlags = topologicalSort(flags, flagKeys);
+    const results = this.evaluation.evaluate(context, sortedFlags);
+    void this.assignmentService?.track(new Assignment(user, results));
+    this.logger.debug('[Experiment] evaluate - variants: ', results);
+    return evaluationVariantsToVariants(results);
+  }
+
   /**
    * Locally evaluates flag variants for a user.
    *
@@ -105,41 +138,14 @@ export class LocalEvaluationClient {
    * @param flagKeys The flags to evaluate with the user. If empty, all flags
    * from the flag cache are evaluated.
    * @returns The evaluated variants
+   * @deprecated use evaluateV2 instead
    */
   public async evaluate(
     user: ExperimentUser,
     flagKeys?: string[],
   ): Promise<Variants> {
-    this.logger.debug(
-      '[Experiment] evaluate - user:',
-      user,
-      'flags:',
-      this.flags,
-    );
-    const results: Results = evaluation.evaluate(this.flags, user);
-    const assignmentResults: Results = {};
-    const variants: Variants = {};
-    const filter = flagKeys && flagKeys.length > 0;
-    for (const flagKey in results) {
-      const included = !filter || flagKeys.includes(flagKey);
-      if (included) {
-        const flagResult = results[flagKey];
-        variants[flagKey] = {
-          value: flagResult.value,
-          payload: flagResult.payload,
-        };
-      }
-      if (
-        included ||
-        results[flagKey].type == FLAG_TYPE_MUTUAL_EXCLUSION_GROUP ||
-        results[flagKey].type == FLAG_TYPE_HOLDOUT_GROUP
-      ) {
-        assignmentResults[flagKey] = results[flagKey];
-      }
-    }
-    void this.assignmentService?.track(new Assignment(user, assignmentResults));
-    this.logger.debug('[Experiment] evaluate - variants: ', variants);
-    return variants;
+    const results = this.evaluateV2(user, flagKeys);
+    return filterDefaultVariants(results);
   }
 
   /**