Fix to allow disabling telemetry

SciPhi-AI · Sep 25, 2024 · 6323145 · 6323145
1 parent 6cf898b
commit 6323145
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 29 deletions.
diff --git a/js/sdk/src/feature.ts b/js/sdk/src/feature.ts
@@ -15,11 +15,16 @@ function isBrowser(
   return typeof window !== "undefined";
 }
 
-export function initializeTelemetry() {
+let telemetryEnabled = true;
+
+export function initializeTelemetry(enableTelemetry: boolean = true) {
   if (isPostHogInitialized) {
     return;
   }
 
+  telemetryEnabled =
+    enableTelemetry && process.env.R2R_JS_DISABLE_TELEMETRY !== "true";
+
   if (typeof window !== "undefined") {
     // Browser environment
     import("posthog-js").then((posthogJs) => {
@@ -28,31 +33,30 @@ export function initializeTelemetry() {
         api_host: posthogHost,
       });
 
+      if (!telemetryEnabled) {
+        posthog.opt_out_capturing();
+      }
+
       window.addEventListener("beforeunload", () => {
-        (posthog as PostHogJs.PostHog).capture("PageUnload");
+        if (telemetryEnabled) {
+          (posthog as PostHogJs.PostHog).capture("PageUnload");
+        }
       });
 
       isPostHogInitialized = true;
     });
   } else {
     // Node.js environment
     posthog = new PostHogNodeClient(posthogApiKey, { host: posthogHost });
-    isPostHogInitialized = true;
-  }
-
-  if (process.env.R2R_JS_DISABLE_TELEMETRY === "true") {
-    if (isBrowser(posthog)) {
-      posthog.opt_out_capturing();
-    } else {
+    if (!telemetryEnabled) {
       posthog.disable();
     }
+    isPostHogInitialized = true;
   }
 }
 
-type AsyncFunction = (...args: any[]) => Promise<any>;
-
 function captureEvent(eventName: string, properties?: Record<string, any>) {
-  if (isPostHogInitialized) {
+  if (isPostHogInitialized && telemetryEnabled) {
     const environment = typeof window !== "undefined" ? "browser" : "node";
     const eventProperties = { ...properties, environment };
 
@@ -68,6 +72,8 @@ function captureEvent(eventName: string, properties?: Record<string, any>) {
   }
 }
 
+type AsyncFunction = (...args: any[]) => Promise<any>;
+
 export function feature(operationName: string) {
   return function (
     _target: any,
@@ -82,17 +88,21 @@ export function feature(operationName: string) {
     ): Promise<any> {
       try {
         const result = await originalMethod.apply(this, args);
-        captureEvent("TSClient", { operation: operationName });
+        if (this.anonymousTelemetry) {
+          captureEvent("TSClient", { operation: operationName });
+        }
         return result;
       } catch (error: unknown) {
-        captureEvent("TSClient", {
-          operation: operationName,
-          errorMessage:
-            error instanceof Error ? error.message : "Unknown error",
-        });
+        if (this.anonymousTelemetry) {
+          captureEvent("TSClient", {
+            operation: operationName,
+            errorMessage:
+              error instanceof Error ? error.message : "Unknown error",
+          });
+        }
         throw error;
       } finally {
-        if (isPostHogInitialized && !isBrowser(posthog)) {
+        if (isPostHogInitialized && !isBrowser(posthog) && telemetryEnabled) {
           // Flush events in Node.js environment
           await (posthog as PostHogNodeClient).shutdown();
         }
@@ -124,16 +134,20 @@ export function featureGenerator(operationName: string) {
         for await (const chunk of generator) {
           yield chunk;
         }
-        captureEvent("TSClient", { operation: operationName });
+        if (this.anonymousTelemetry) {
+          captureEvent("TSClient", { operation: operationName });
+        }
       } catch (error: unknown) {
-        captureEvent("TSClient", {
-          operation: operationName,
-          errorMessage:
-            error instanceof Error ? error.message : "Unknown error",
-        });
+        if (this.anonymousTelemetry) {
+          captureEvent("TSClient", {
+            operation: operationName,
+            errorMessage:
+              error instanceof Error ? error.message : "Unknown error",
+          });
+        }
         throw error;
       } finally {
-        if (isPostHogInitialized && !isBrowser(posthog)) {
+        if (isPostHogInitialized && !isBrowser(posthog) && telemetryEnabled) {
           // Flush events in Node.js environment
           await (posthog as PostHogNodeClient).shutdown();
         }

diff --git a/js/sdk/src/r2rClient.ts b/js/sdk/src/r2rClient.ts
@@ -53,14 +53,25 @@ function handleRequestError(response: AxiosResponse): void {
 export class r2rClient {
   private axiosInstance: AxiosInstance;
   private baseUrl: string;
+  private anonymousTelemetry: boolean;
+
+  // Authorization tokens
   private accessToken: string | null;
   private refreshToken: string | null;
 
-  constructor(baseURL: string, prefix: string = "/v2") {
+  constructor(
+    baseURL: string,
+    prefix: string = "/v2",
+    anonymousTelemetry = true,
+  ) {
     this.baseUrl = `${baseURL}${prefix}`;
+    this.anonymousTelemetry = anonymousTelemetry;
+
     this.accessToken = null;
     this.refreshToken = null;
 
+    initializeTelemetry(this.anonymousTelemetry);
+
     this.axiosInstance = axios.create({
       baseURL: this.baseUrl,
       headers: {
@@ -90,8 +101,6 @@ export class r2rClient {
         },
       ],
     });
-
-    initializeTelemetry();
   }
 
   setTokens(accessToken: string, refreshToken: string): void {