(frontend)[AGE-1418]: Improve posthog initialization to reduce load on bundle size

agenta-web/src/components/Playground/Views/ParametersView.tsx

@@ -7,7 +7,7 @@ import {ModelParameters, ObjectParameters, StringParameters} from "./ParametersC
 import PublishVariantModal from "./PublishVariantModal"
 import {deleteSingleVariant} from "@/services/playground/api"
 import {CloudUploadOutlined, DeleteOutlined, HistoryOutlined, SaveOutlined} from "@ant-design/icons"
-import {usePostHogAg} from "@/hooks/usePostHogAg"
+import {usePostHogAg} from "@/lib/helpers/analytics/hooks/usePostHogAg"
 import {isDemo} from "@/lib/helpers/utils"
 import {useQueryParam} from "@/hooks/useQuery"
 import {dynamicComponent, dynamicService} from "@/lib/helpers/dynamic"
@@ -128,7 +128,7 @@ const ParametersView: React.FC<Props> = ({
                 onStateChange(false)
                 res(true)
             })
-            posthog.capture("variant_saved", {variant_id: variant.variantId})
+            posthog?.capture?.("variant_saved", {variant_id: variant.variantId})
         })
     }
 

agenta-web/src/components/Playground/Views/PublishVariantModal.tsx

@@ -1,4 +1,4 @@
-import {usePostHogAg} from "@/hooks/usePostHogAg"
+import {usePostHogAg} from "@/lib/helpers/analytics/hooks/usePostHogAg"
 import {Environment, Variant} from "@/lib/Types"
 import {variantNameWithRev} from "@/lib/helpers/variantHelper"
 import {fetchEnvironments, createPublishVariant} from "@/services/deployment/api"
@@ -58,7 +58,7 @@ const PublishVariantModal: React.FC<Props> = ({
             closeModal()
             await loadEnvironments()
             message.success(`Published ${variant.variantName} to ${envName}`)
-            posthog.capture("app_deployed", {app_id: appId, environment: envName})
+            posthog?.capture?.("app_deployed", {app_id: appId, environment: envName})
         })
     }
 

agenta-web/src/components/pages/app-management/index.tsx

@@ -7,7 +7,7 @@ import {waitForAppToStart} from "@/services/api"
 import {createUseStyles} from "react-jss"
 import {useAppsData} from "@/contexts/app.context"
 import {useProfileData} from "@/contexts/profile.context"
-import {usePostHogAg} from "@/hooks/usePostHogAg"
+import {usePostHogAg} from "@/lib/helpers/analytics/hooks/usePostHogAg"
 import {LlmProvider, getAllProviderLlmKeys} from "@/lib/helpers/llmProviders"
 import {dynamicComponent, dynamicContext} from "@/lib/helpers/dynamic"
 import dayjs from "dayjs"
@@ -130,7 +130,7 @@ const AppManagement: React.FC = () => {
                     setFetchingTemplate(false)
                 if (status === "success") {
                     mutate()
-                    posthog.capture("app_deployment", {
+                    posthog?.capture?.("app_deployment", {
                         properties: {
                             app_id: appId,
                             environment: "UI",

agenta-web/src/contexts/profile.context.tsx

@@ -1,4 +1,4 @@
-import {usePostHogAg} from "@/hooks/usePostHogAg"
+import {usePostHogAg} from "@/lib/helpers/analytics/hooks/usePostHogAg"
 import {useSession} from "@/hooks/useSession"
 import useStateCallback from "@/hooks/useStateCallback"
 import {isDemo} from "@/lib/helpers/utils"
@@ -38,7 +38,7 @@ const ProfileContextProvider: React.FC<PropsWithChildren> = ({children}) => {
         setLoading(true)
         fetchProfile()
             .then((profile) => {
-                posthog.identify()
+                posthog?.identify?.()
                 setUser(profile.data, onSuccess)
             })
             .catch((error) => {

agenta-web/src/hooks/useSession.ts

@@ -1,7 +1,6 @@
 import {useProfileData} from "@/contexts/profile.context"
 import {isDemo} from "@/lib/helpers/utils"
 import {useRouter} from "next/router"
-import posthog from "posthog-js"
 import {useSessionContext} from "supertokens-auth-react/recipe/session"
 import {signOut} from "supertokens-auth-react/recipe/session"
 
@@ -17,7 +16,8 @@ export const useSession: () => {loading: boolean; doesSessionExist: boolean; log
                   doesSessionExist: (res as any).doesSessionExist,
                   logout: () => {
                       signOut()
-                          .then(() => {
+                          .then(async () => {
+                              const posthog = (await import("posthog-js")).default
                               posthog.reset()
                               reset()
                               router.push("/auth")

agenta-web/src/lib/helpers/analytics/AgPosthogProvider.tsx

@@ -0,0 +1,55 @@
+import {useCallback, useEffect, useRef, useState} from "react"
+import {useRouter} from "next/router"
+import {useAtom} from "jotai"
+import {posthogAtom, type PostHogConfig} from "./store/atoms"
+import {CustomPosthogProviderType} from "./types"
+
+const CustomPosthogProvider: CustomPosthogProviderType = ({children, config}) => {
+    const router = useRouter()
+    const [loadingPosthog, setLoadingPosthog] = useState(false)
+    const [posthogClient, setPosthogClient] = useAtom(posthogAtom)
+
+    const initPosthog = useCallback(async () => {
+        if (!!posthogClient) return
+        if (loadingPosthog) return
+
+        setLoadingPosthog(true)
+
+        try {
+            const posthog = (await import("posthog-js")).default
+
+            posthog.init(process.env.NEXT_PUBLIC_POSTHOG_API_KEY!, {
+                api_host: "https://app.posthog.com",
+                // Enable debug mode in development
+                loaded: (posthog) => {
+                    setPosthogClient(posthog)
+                    if (process.env.NODE_ENV === "development") posthog.debug()
+                },
+                capture_pageview: false,
+                ...config,
+            })
+        } finally {
+            setLoadingPosthog(false)
+        }
+    }, [loadingPosthog, config, posthogClient, setPosthogClient])
+
+    useEffect(() => {
+        initPosthog()
+    }, [initPosthog])
+
+    const handleRouteChange = useCallback(() => {
+        posthogClient?.capture("$pageview", {$current_url: window.location.href})
+    }, [posthogClient])
+
+    useEffect(() => {
+        router.events.on("routeChangeComplete", handleRouteChange)
+
+        return () => {
+            router.events.off("routeChangeComplete", handleRouteChange)
+        }
+    }, [handleRouteChange, router.events])
+
+    return <>{children}</>
+}
+
+export default CustomPosthogProvider

agenta-web/src/lib/helpers/analytics/hooks/usePostHogAg.ts

@@ -0,0 +1,51 @@
+import {isDemo, generateOrRetrieveDistinctId} from "@/lib/helpers/utils"
+import {useProfileData} from "@/contexts/profile.context"
+import {useAtom} from "jotai"
+import {posthogAtom} from "../store/atoms"
+import {type PostHog} from "posthog-js"
+import useIsomorphicLayoutEffect from "@/hooks/useIsomorphicLayoutEffect"
+
+interface ExtendedPostHog extends PostHog {
+    identify: PostHog["identify"]
+    capture: PostHog["capture"]
+}
+
+export const usePostHogAg = (): ExtendedPostHog | null => {
+    const trackingEnabled = process.env.NEXT_PUBLIC_TELEMETRY_TRACKING_ENABLED === "true"
+    const {user} = useProfileData()
+    const [posthog] = useAtom(posthogAtom)
+
+    const _id: string | undefined = isDemo() ? user?.email : generateOrRetrieveDistinctId()
+    const capture: PostHog["capture"] = (...args) => {
+        if (trackingEnabled && user?.id) {
+            return posthog?.capture?.(...args)
+        }
+        return undefined
+    }
+    const identify: PostHog["identify"] = (id, ...args) => {
+        if (trackingEnabled && user?.id) {
+            posthog?.identify?.(_id !== undefined ? _id : id, ...args)
+        }
+    }
+    useIsomorphicLayoutEffect(() => {
+        if (!posthog) return
+
+        if (!trackingEnabled) {
+            console.log("POSTHOG: opt_out_capturing")
+            posthog.opt_out_capturing()
+        }
+    }, [posthog, trackingEnabled])
+
+    useIsomorphicLayoutEffect(() => {
+        if (!posthog) return
+        if (posthog.get_distinct_id() !== _id) identify()
+    }, [posthog, _id])
+
+    return posthog
+        ? ({
+              ...posthog,
+              identify,
+              capture,
+          } as ExtendedPostHog)
+        : null
+}

agenta-web/src/lib/helpers/analytics/store/atoms.ts

@@ -0,0 +1,5 @@
+import {atom} from "jotai"
+import {type PostHog, type PostHogConfig} from "posthog-js"
+
+export type {PostHogConfig}
+export const posthogAtom = atom<PostHog | null>(null)

agenta-web/src/lib/helpers/analytics/types.d.ts

@@ -0,0 +1,7 @@
+import {type PostHogConfig} from "./store/atoms"
+
+export interface CustomPosthogProviderType
+    extends React.FC<{
+        children: React.ReactNode
+        config: Partial<PostHogConfig>
+    }> {}

agenta-web/src/pages/_app.tsx

@@ -1,12 +1,7 @@
-import {useEffect} from "react"
 import type {AppProps} from "next/app"
-import {useRouter} from "next/router"
 import Head from "next/head"
 import dynamic from "next/dynamic"
 
-import posthog from "posthog-js"
-import {PostHogProvider} from "posthog-js/react"
-
 import "@/styles/globals.css"
 import Layout from "@/components/Layout/Layout"
 import {dynamicComponent} from "@/lib/helpers/dynamic"
@@ -19,45 +14,26 @@ import "ag-grid-community/styles/ag-theme-alpine.css"
 import {Inter} from "next/font/google"
 
 const NoMobilePageWrapper = dynamicComponent("NoMobilePageWrapper/NoMobilePageWrapper")
+const CustomPosthogProvider = dynamic(() => import("@/lib/helpers/analytics/AgPosthogProvider"))
 
 const inter = Inter({
     subsets: ["latin"],
     variable: "--font-inter",
 })
 
-// Initialize the Posthog client
-if (typeof window !== "undefined") {
-    posthog.init(process.env.NEXT_PUBLIC_POSTHOG_API_KEY!, {
-        api_host: "https://app.posthog.com",
-        // Enable debug mode in development
-        loaded: (posthog) => {
-            if (process.env.NODE_ENV === "development") posthog.debug()
-        },
-        capture_pageview: false,
-        persistence: "localStorage+cookie",
-    })
-}
-
 export default function App({Component, pageProps}: AppProps) {
-    const router = useRouter()
-
-    useEffect(() => {
-        const handleRouteChange = () =>
-            posthog.capture("$pageview", {$current_url: window.location.href})
-        router.events.on("routeChangeComplete", handleRouteChange)
-
-        return () => {
-            router.events.off("routeChangeComplete", handleRouteChange)
-        }
-    }, [])
     return (
         <>
             <Head>
                 <title>Agenta: The LLMOps platform.</title>
                 <link rel="shortcut icon" href="/assets/favicon.ico" />
             </Head>
             <main className={`${inter.variable} font-sans`}>
-                <PostHogProvider client={posthog}>
+                <CustomPosthogProvider
+                    config={{
+                        persistence: "localStorage+cookie",
+                    }}
+                >
                     <ThemeContextProvider>
                         <ProfileContextProvider>
                             <ProjectContextProvider>
@@ -70,7 +46,7 @@ export default function App({Component, pageProps}: AppProps) {
                             </ProjectContextProvider>
                         </ProfileContextProvider>
                     </ThemeContextProvider>
-                </PostHogProvider>
+                </CustomPosthogProvider>
             </main>
         </>
     )